Update lxmert/src/ExplanationGenerator.py

lxmert/src/ExplanationGenerator.py

@@ -163,7 +163,7 @@ class GeneratorOurs:
         one_hot[0, index] = 1
         one_hot_vector = one_hot
         one_hot = torch.from_numpy(one_hot).requires_grad_(True)
-        one_hot = torch.sum(one_hot.cuda() * output)
+        one_hot = torch.sum(one_hot * output)
 
         model.zero_grad()
         one_hot.backward(retain_graph=True)
@@ -400,7 +400,7 @@ class GeneratorBaselines:
         one_hot[0, index] = 1
         one_hot_vector = one_hot
         one_hot = torch.from_numpy(one_hot).requires_grad_(True)
-        one_hot = torch.sum(one_hot.cuda() * output)
+        one_hot = torch.sum(one_hot * output)
 
         model.zero_grad()
         one_hot.backward(retain_graph=True)