fix import

AI4SIM · Jan 29, 2025 · 8271deb · 8271deb
1 parent 26d8915
commit 8271deb
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/reactive-flows/cnf-combustion/gnns/tests/test_models.py b/reactive-flows/cnf-combustion/gnns/tests/test_models.py
@@ -23,9 +23,9 @@
 import torch_geometric as pyg
 import torch_optimizer as optim
 import yaml
-
 from lightning.pytorch.trainer import Trainer
-from models import LitGCN, LitGAT, LitGIN, LitGraphUNet
+
+from models import LitGAT, LitGCN, LitGIN, LitGraphUNet
 
 
 class TestModel(unittest.TestCase):
@@ -171,9 +171,7 @@ def test_test_step(self):
             test_gunet.on_test_epoch_end()
 
             self.assertTrue(
-                os.path.exists(
-                    os.path.join(test_gunet.trainer.log_dir, "plots")
-                )
+                os.path.exists(os.path.join(test_gunet.trainer.log_dir, "plots"))
             )
 
     def test_configure_optimizers(self):