Update lxmert/src/ExplanationGenerator.py

lxmert/src/ExplanationGenerator.py

@@ -317,7 +317,7 @@ class GeneratorOursAblationNoAggregation:
         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)