Update test Model construction.

choderalab · Nov 29, 2023 · 4061e2e · 4061e2e
1 parent 0c39b98
commit 4061e2e
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/mtenn/tests/test_combination.py b/mtenn/tests/test_combination.py
@@ -2,7 +2,6 @@
 import numpy as np
 import pytest
 import torch
-from torch_geometric.nn import SchNet as PygSchNet
 
 from mtenn.combination import MeanCombination, MaxCombination, BoltzmannCombination
 from mtenn.conversion_utils import SchNet
@@ -11,7 +10,7 @@
 @pytest.fixture()
 def models_and_inputs():
     model_test = SchNet(
-        PygSchNet(hidden_channels=2, num_filters=2, num_interactions=2, num_gaussians=2)
+        hidden_channels=2, num_filters=2, num_interactions=2, num_gaussians=2
     )
     model_ref = deepcopy(model_test)
     model_ref = SchNet.get_model(model_ref, strategy="complex")