Removed distributed training and fixed issue with Tag name not being …

…used as run name
GeroVanMi · May 15, 2024 · 0c54ae5 · 0c54ae5
1 parent 9d87d1f
commit 0c54ae5
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/ci/start_lightning_job.py b/ci/start_lightning_job.py
@@ -8,7 +8,7 @@
     studio = create_studio()
 
     print("Starting studio.")
-    studio.start(Machine.T4_X_4)
+    studio.start(Machine.T4)
 
     BASE_PATH = "/teamspace/studios/this_studio/algorithmic-quartet-mlops"
 
@@ -28,5 +28,5 @@
     )
 
     studio.run(
-        f"docker run -d --gpus all -e TAG_NAME -e WANDB_API_KEY='{os.environ.get('WANDB_API_KEY')}' -e GC_BUCKET_KEY='{os.environ.get('GC_BUCKET_KEY')}' us-west2-docker.pkg.dev/algorithmic-quartet/training-images/pokemon-trainer:latest"
+        f"docker run -d --gpus all -e TAG_NAME={os.environ.get('TAG_NAME')} -e WANDB_API_KEY='{os.environ.get('WANDB_API_KEY')}' -e GC_BUCKET_KEY='{os.environ.get('GC_BUCKET_KEY')}' us-west2-docker.pkg.dev/algorithmic-quartet/training-images/pokemon-trainer:latest"
     )
diff --git a/training/src/configurations/TrainConfig.py b/training/src/configurations/TrainConfig.py
@@ -5,7 +5,7 @@
 
 @dataclass
 class TrainConfig(Configuration):
-    train_batch_size = 96
+    train_batch_size = 24
     eval_batch_size = 8
     num_epochs = 60
     num_images = 7351