Upload dti.py

deepscreen/models/dti.py

@@ -1,5 +1,5 @@
 from functools import partial
-from typing import Optional, Sequence, Dict
+from typing import Optional, Sequence, Dict, Any
 
 from torch import nn, optim, Tensor
 from lightning import LightningModule
@@ -53,12 +53,12 @@ class DTILightningModule(LightningModule):
                 dataloader = self.trainer.datamodule.train_dataloader()
                 dummy_batch = next(iter(dataloader))
                 self.forward(dummy_batch)
-            case 'validate':
-                dataloader = self.trainer.datamodule.val_dataloader()
-            case 'test':
-                dataloader = self.trainer.datamodule.test_dataloader()
-            case 'predict':
-                dataloader = self.trainer.datamodule.predict_dataloader()
+            # case 'validate':
+            #     dataloader = self.trainer.datamodule.val_dataloader()
+            # case 'test':
+            #     dataloader = self.trainer.datamodule.test_dataloader()
+            # case 'predict':
+            #     dataloader = self.trainer.datamodule.predict_dataloader()
 
         # for key, value in dummy_batch.items():
         #     if isinstance(value, Tensor):