Update DETR/modules/ExplanationGenerator.py

DETR/modules/ExplanationGenerator.py

@@ -157,7 +157,7 @@ class Generator:
         one_hot[0, target_index, index] = 1
         one_hot_vector = one_hot
         one_hot.requires_grad_(True)
-        one_hot = torch.sum(one_hot.cuda() * outputs)
+        one_hot = torch.sum(one_hot * outputs)
 
         self.model.zero_grad()
         one_hot.backward(retain_graph=True)