fix imports in train.py (microsoft#332)

yichiac · Dec 27, 2021 · c6f81a4 · c6f81a4
1 parent c1bdbde
commit c6f81a4
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/train.py b/train.py
@@ -13,7 +13,7 @@
 from pytorch_lightning import loggers as pl_loggers
 from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
 
-from .datamodules import (
+from torchgeo.datamodules import (
     BigEarthNetDataModule,
     ChesapeakeCVPRDataModule,
     COWCCountingDataModule,
@@ -28,17 +28,17 @@
     So2SatDataModule,
     UCMercedDataModule,
 )
-from .trainers import (
+from torchgeo.trainers import (
     BYOLTask,
     ClassificationTask,
     MultiLabelClassificationTask,
     RegressionTask,
     SemanticSegmentationTask,
 )
-from .trainers.chesapeake import ChesapeakeCVPRSegmentationTask
-from .trainers.landcoverai import LandCoverAISegmentationTask
-from .trainers.naipchesapeake import NAIPChesapeakeSegmentationTask
-from .trainers.resisc45 import RESISC45ClassificationTask
+from torchgeo.trainers.chesapeake import ChesapeakeCVPRSegmentationTask
+from torchgeo.trainers.landcoverai import LandCoverAISegmentationTask
+from torchgeo.trainers.naipchesapeake import NAIPChesapeakeSegmentationTask
+from torchgeo.trainers.resisc45 import RESISC45ClassificationTask
 
 TASK_TO_MODULES_MAPPING: Dict[
     str, Tuple[Type[pl.LightningModule], Type[pl.LightningDataModule]]