fix

kumo-ai · Oct 1, 2024 · 03160d8 · 03160d8
1 parent 8107b2c
commit 03160d8
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/hybridgnn/nn/models/wmf.py b/hybridgnn/nn/models/wmf.py
@@ -46,8 +46,5 @@ def forward(
         neg_lhs = self.full_lhs[mask]
         mask = ~torch.isin(self.full_rhs, dst_tensor)
         neg_rhs = self.full_rhs[mask]
-        from torch.cuda.amp import autocast
-
-        with autocast():
-            mat_neg = torch.mm(self.lhs(neg_lhs).half(), self.rhs(neg_rhs).half().t())
+        mat_neg = torch.mm(self.lhs(neg_lhs), self.rhs(neg_rhs).t())
         return ((1.0 - mat_pos) **2).sum() + self.w0*((mat_neg**2).sum())