add drop_last in mnist eager example

nod-ai · Oct 14, 2023 · 347a561 · 347a561
1 parent 3e35a39
commit 347a561
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/eager_mlp/mlp_eager_simple.py b/examples/eager_mlp/mlp_eager_simple.py
@@ -39,7 +39,7 @@ def get_train_loader(self):
 
     def get_test_loader(self):
         return DataLoader(
-            dataset=self.mnist_testset, batch_size=self.batch_size, shuffle=False
+            dataset=self.mnist_testset, batch_size=self.batch_size, shuffle=False, drop_last=True,
         )
 
 
@@ -70,7 +70,7 @@ def infer_iteration(model, images):
 def infer():
     # Example Parameters
     config = {
-        "batch_size": 100,
+        "batch_size": 64,
         "learning_rate": 0.001,
         "num_epochs": 10,
     }