Merge pull request #21 from gpu-mode/triton_check

Re-order write to train.py and install accordingly

.github/workflows/train_workflow.yml

@@ -19,17 +19,29 @@ jobs:
       with:
         python-version: '3.10'
 
-    - name: Install dependencies
-      run: |
-        pip install numpy
-        pip install torch
-
     - name: Create training script
       shell: python
       run: |
         with open('train.py', 'w') as f:
             f.write('''${{ github.event.inputs.script_content }}''')
 
+    - name: Install dependencies
+      run: |
+        if grep -rE "(import numpy|from numpy)" train.py; then
+          echo "Numpy detected, installing numpy"
+          pip install numpy
+        fi
+        # Check if 'import torch' is in any Python file
+        if grep -rE "(import torch|from torch)" train.py; then
+          echo "PyTorch detected, installing torch"
+          pip install torch
+        fi
+        # Check if 'import triton' is in any Python file
+        if grep -rE "(import triton|from triton)" train.py; then
+          echo "Triton detected, installing triton"
+          pip install triton
+        fi
+
     - name: Run training script
       run: |
         python train.py > training.log 2>&1