update trainer

openforcefield · Feb 9, 2024 · e8a70d6 · e8a70d6
1 parent 922a314
commit e8a70d6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/openff/nagl/tests/training/test_training.py b/openff/nagl/tests/training/test_training.py
@@ -340,5 +340,5 @@ def test_weighted_mixed_training_step(self, mock_training_model, dgl_methane):
 def test_train_model_no_error(example_training_config):
     model = TrainingGNNModel(example_training_config)
     data = DGLMoleculeDataModule(example_training_config)
-    trainer = pl.Trainer(gpus=1, max_epochs=2)
+    trainer = pl.Trainer(accelerator="gpu", devices=1, max_epochs=2)
     trainer.fit(model, datamodule=data)