minor bug fix

ninadakolekar · Dec 6, 2018 · 8a2de4e · 8a2de4e
1 parent 681640f
commit 8a2de4e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/models/ian.py b/models/ian.py
@@ -38,8 +38,10 @@ def forward(self, inputs):
         context_pool = torch.sum(context, dim=1)
         context_pool = torch.div(context_pool, text_raw_len.view(text_raw_len.size(0), 1))
 
-        aspect_final, _ = self.attention_aspect(aspect, context_pool).squeeze(dim=1)
-        context_final, _ = self.attention_context(context, aspect_pool).squeeze(dim=1)
+        aspect_final, _ = self.attention_aspect(aspect, context_pool)
+        aspect_final = aspect_final.squeeze(dim=1)
+        context_final, _ = self.attention_context(context, aspect_pool)
+        context_final = context_final.squeeze(dim=1)
 
         x = torch.cat((aspect_final, context_final), dim=-1)
         out = self.dense(x)

diff --git a/train.py b/train.py
@@ -24,7 +24,8 @@ def __init__(self, opt):
         self.test_data_loader = DataLoader(dataset=absa_dataset.test_data, batch_size=opt.batch_size, shuffle=False)
 
         self.model = opt.model_class(absa_dataset.embedding_matrix, opt).to(opt.device)
-        print("cuda memory allocated:", torch.cuda.memory_allocated(device=0))
+        if opt.device.type == 'cuda':
+            print("cuda memory allocated:", torch.cuda.memory_allocated(device=opt.device.index))
         self._print_args()
 
     def _print_args(self):