diff --git a/.github/workflows/train_workflow.yml b/.github/workflows/train_workflow.yml
index 53d55c2..eae9c9a 100644
--- a/.github/workflows/train_workflow.yml
+++ b/.github/workflows/train_workflow.yml
@@ -9,12 +9,11 @@ on:
 
 jobs:
   train:
-    runs-on: ubuntu-latest
+    runs-on: [gpumode-nvidia-arc, amd-mi250] # Add more GPUs as we get them here
     steps:
       - name: Install dependencies
         run: |
-          pip install numpy
-          # pip install torch - need to find a way to cache this otherwise it will take a long time to install
+          pip install numpy torch
 
       - name: Create and run training script
         run: |
diff --git a/train.py b/train.py
index cd867ab..684f3e6 100644
--- a/train.py
+++ b/train.py
@@ -1,8 +1,12 @@
-import numpy
+import torch
 
-a = numpy.array([1, 2, 3])
-b = numpy.array([4, 5, 6])
 
+
+# Vector addition
+a = torch.tensor([1, 2, 3]).cuda()
+b = torch.tensor([4, 5, 6]).cuda()
 c = a + b
 
-print(c)
\ No newline at end of file
+print(a)
+print(b)
+print(c)