changed dependencies

Dockerfile

@@ -13,8 +13,6 @@ RUN pip install --no-cache-dir torch torchvision
 COPY gpu_check.py /code/gpu_check.py
 RUN python /code/gpu_check.py
 
-RUN nvcc -V
-
 # Set up a new user named "user" with user ID 1000
 RUN useradd -m -u 1000 user
 

main.py

@@ -37,6 +37,15 @@ print(pynvml.nvmlSystemGetCudaDriverVersion())
 cuda_version = torch.version.cuda
 print("PyTorch is using CUDA Version:", cuda_version)
 
+# Check if CUDA is available
+if torch.cuda.is_available():
+    print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+    device = torch.device("cuda")
+else:
+    print("CUDA is not available. Using CPU.")
+
+torch.zeros(1).cuda()
+
 
 def print_gpu_utilization():
     nvmlInit()
@@ -196,16 +205,6 @@ def initModel(model_name_or_path, revision):
     return generator
 
 
-# Check if CUDA is available
-if torch.cuda.is_available():
-    print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
-    device = torch.device("cuda")
-else:
-    print("CUDA is not available. Using CPU.")
-
-torch.zeros(1).cuda()
-
-
 def generateTags():
     start = time.time()
     res = generateMods(generator, bag_metadata, dress_example, bomber_example)