From 4061e2eda04ee1ded02c6359299c592701ff56f8 Mon Sep 17 00:00:00 2001 From: kaminow Date: Wed, 29 Nov 2023 11:39:02 -0500 Subject: [PATCH] Update test Model construction. --- mtenn/tests/test_combination.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mtenn/tests/test_combination.py b/mtenn/tests/test_combination.py index 13355c3..ac57fd8 100644 --- 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")