Update test_main.py

DLR-MF-DAS · Feb 21, 2025 · cebf357 · cebf357
1 parent 030add7
commit cebf357
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/test/test_main.py b/test/test_main.py
@@ -4,8 +4,8 @@
 from methane_super_emitters.datamodule import TROPOMISuperEmitterDataModule
 
 def test_model():
-    model = SuperEmitterDetector()
-    datamodule = TROPOMISuperEmitterDataModule('./data/dataset_wf')
-    trainer = L.Trainer(max_epochs=1)
+    model = SuperEmitterDetector(fields=['methane', 'qa', 'u10', 'v10'])
+    datamodule = TROPOMISuperEmitterDataModule('./data/dataset', fields=['methane', 'qa', 'u10', 'v10'])
+    trainer = L.Trainer(max_epochs=10)
     trainer.fit(model=model, datamodule=datamodule)
     trainer.test(model=model, datamodule=datamodule)