NVIDA test

.github/workflows/train_workflow.yml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   train:
-    runs-on: [amdgpu-mi250-x86-64]
+    runs-on: [gpumode-nvidia-arc]
 
     steps:
     - name: Setup Python
@@ -21,7 +21,7 @@ jobs:
     - name: Install dependencies
       run: |
         pip install numpy
-        pip install torch --index-url https://download.pytorch.org/whl/rocm6.2
+        pip install torch
 
     - name: Create training script
       shell: python

train.py

@@ -1,12 +1,8 @@
 import torch
 
-a = torch.Tensor([1, 2, 3, 4, 5]).to("cuda")
-b= torch.Tensor([1, 2, 3, 4, 5]).to("cuda")
-
-if torch.cuda.is_available():
-    gpu_name = torch.cuda.get_device_name(0)
-    print(f"GPU Name: {gpu_name}")
-else:
-    print("No GPU available")
+a = torch.Tensor([1, 2, 3, 4, 5]).cuda()
+b= torch.Tensor([1, 2, 3, 4, 5]).cuda()
 
+print(a)
+print(b)
 print(a + b)