bs 4 -> 2 in training test
MODIFY
test/test_train.py
test/test_train.py
@@ -11,7 +11,7 @@ from models.transformer import Transformer
from models.vit import ViT
from models.resnet import ResNet18
-BS = int(os.getenv("BS", "4"))
+BS = int(os.getenv("BS", "2"))
def train_one_step(model,X,Y):
params = get_parameters(model)

bs 4 -> 2 in training test