update

pyg-team · Nov 1, 2024 · b4a3c2d · b4a3c2d
1 parent 0012833
commit b4a3c2d
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/examples/llm/git_mol.py b/examples/llm/git_mol.py
@@ -66,8 +66,6 @@ def train(
     test_loader = accelerator.prepare_data_loader(test_loader,
                                                   device_placement=True)
 
-    # import pdb
-    # pdb.set_trace()
     # Train and eval ============================================
     best_epoch = 0
     best_val_loss = float('inf')

diff --git a/test/nn/models/test_git_mol.py b/test/nn/models/test_git_mol.py
@@ -8,12 +8,12 @@
 def test_git_mol():
     model = GITMol()
 
-    x = torch.randn(10, 16)
+    x = torch.ones(10, 16, dtype=torch.long)
     edge_index = torch.tensor([
         [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
         [1, 2, 3, 4, 0, 6, 7, 8, 9, 5],
     ])
-    edge_attr = torch.randn(edge_index.size(1), 16)
+    edge_attr = torch.zeros(edge_index.size(1), 16, dtype=torch.long)
     batch = torch.tensor([0, 0, 0, 0, 0, 1, 1, 1, 1, 1])
     smiles = ['CC(C)([C@H]1CC2=C(O1)C=CC3=C2OC(=O)C=C3)O'] * 2
     captions = ['The molecule is the (R)-(-)-enantiomer of columbianetin.'] * 2