FRML-67 Fix CML not leveraging GPU

.github/workflows/model.yml

@@ -12,7 +12,9 @@ jobs:
       volumes:
         - /home/runner/work/frdc-ml/_github_home:/root
       env:
+        # This is where setup-python will install and cache the venv
         AGENT_TOOLSDIRECTORY: "/root/venv"
+      options: --gpus all
 
     steps:
       - uses: actions/checkout@v3
@@ -34,7 +36,10 @@ jobs:
           pip3 install -r requirements.txt
           pip3 install torch torchvision torchaudio
 
-      - name: Set up gcloud
+      - name: Check CUDA is available
+        run: nvidia-smi
+
+      - name: Auth gcloud
         id: 'auth'
         uses: 'google-github-actions/auth@v1'
         with:
@@ -47,10 +52,12 @@ jobs:
         run: |
           echo "WANDB_API_KEY=${{ secrets.WANDB_API_KEY }}" >> $GITHUB_ENV
 
+      # Our project has src as a source path, explicitly add that in.
       - name: Add src as PYTHONPATH
         run: |
           echo "PYTHONPATH=src" >> $GITHUB_ENV
 
+      # Do not do cd as it'll break PYTHONPATH.
       - name: Run Model Training
         run: |
           python3 -m tests.model_tests.chestnut_dec_may.main

.gitignore

@@ -165,4 +165,5 @@ cython_debug/
 rsc/**/*.tif
 
 **/*/lightning_logs
-*.zip
+*.zip
+*.sh

tests/model_tests/chestnut_dec_may/main.py

@@ -84,7 +84,7 @@ def train_val_test_split(
     # TODO: Though this is set, the results are still not reproducible.
     deterministic=True,
     # fast_dev_run=True,
-    accelerator="cpu",
+    accelerator="gpu",
     log_every_n_steps=4,
     callbacks=[
         # Stop training if the validation loss doesn't improve for 4 epochs