fixes

DSGT-DLP · Nov 19, 2023 · 36501e3 · 36501e3
1 parent 07a70f3
commit 36501e3
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 4 deletions.
diff --git a/dlp-cli b/dlp-cli
diff --git a/training/training/core/dataset.py b/training/training/core/dataset.py
@@ -109,7 +109,7 @@ class DefaultImageDatasets(Enum):
     MNIST = "MNIST"
     FASHION_MNIST = "FashionMNIST"
     KMNIST = "KMNIST"
-    CIFAR = "CIFAR10"
+    CIFAR10 = "CIFAR10"
 
 
 class ImageDefaultDatasetCreator(TrainTestDatasetCreator):

diff --git a/training/training/core/trainer.py b/training/training/core/trainer.py
@@ -204,8 +204,6 @@ def _train_init(self):
         super()._train_init()
 
     def _train_step(self, inputs: torch.Tensor, labels: torch.Tensor):
-        print(f"input shape: {inputs.shape}")
-        print(f"labels shape: {labels.shape}")
         super()._train_step(inputs, labels)
         self._train_correct += self.compute_correct(self._outputs, labels)