From 3e35a39fe4ccc478472756aa222bd66c8337883c Mon Sep 17 00:00:00 2001 From: brucekimrokcmu Date: Sat, 14 Oct 2023 09:48:17 -0700 Subject: [PATCH] add drop_last to ensure all inputs as tensors --- tests/dynamo/mninst_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/dynamo/mninst_test.py b/tests/dynamo/mninst_test.py index 7374dde13..65ad254cd 100644 --- a/tests/dynamo/mninst_test.py +++ b/tests/dynamo/mninst_test.py @@ -44,7 +44,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, ) @@ -75,7 +75,7 @@ def infer_iteration(model, images): def infer(): # Example Parameters config = { - "batch_size": 100, + "batch_size": 64, "learning_rate": 0.001, "num_epochs": 10, }