From dc901d420bdabb849cb573f60d355abd6eb348f7 Mon Sep 17 00:00:00 2001
From: chinthysl <chin.dev.acc@gmail.com>
Date: Wed, 8 May 2024 04:50:24 +0000
Subject: [PATCH] set correct gpu using multigpu config

---
 train_gpt2.cu | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/train_gpt2.cu b/train_gpt2.cu
index 16ff756..15d9158 100644
--- a/train_gpt2.cu
+++ b/train_gpt2.cu
@@ -428,6 +428,7 @@ MultiGpuConfig multi_gpu_config_init(int *argc, char ***argv) {
     return result;
 #else
     printf("Multi-GPU support is disabled. Using a single GPU.\n");
+    cudaCheck(cudaSetDevice(0));
     MultiGpuConfig result;
     result.process_rank = 0;
     result.num_processes = 1;
@@ -2181,12 +2182,10 @@ void gpt2_free(GPT2 *model) {
 // ----------------------------------------------------------------------------
 // common init & free code for train/test/profile
 void common_start(bool override_enable_tf32 = true, bool print_device_info = true) {
-    int deviceIdx = 0;
-    cudaCheck(cudaSetDevice(deviceIdx));
-    cudaGetDeviceProperties(&deviceProp, deviceIdx);
+    cudaGetDeviceProperties(&deviceProp, multi_gpu_config.local_device_idx);
     if (print_device_info) {
         printf("[System]\n");
-        printf("Device %d: %s\n", deviceIdx, deviceProp.name);
+        printf("Device %d: %s\n", multi_gpu_config.local_device_idx, deviceProp.name);
     }
 
     cudaCheck(cudaStreamCreate(&main_stream));