Update resnet.py

resnet.py

@@ -23,6 +23,10 @@ import torchvision.datasets as datasets
 import pytorch_lightning as pl
 import matplotlib.pyplot as plt
 
+
+PATH_DATASETS = os.environ.get("PATH_DATASETS", ".")
+BATCH_SIZE = 256 if AVAIL_GPUS else 64
+
 # Model
 class custom_ResNet(pl.LightningModule):
     def __init__(self, data_dir=PATH_DATASETS, learning_rate=2e-4):