Update inference_from_video.py

inference_from_video.py

@@ -3,6 +3,11 @@ import copy
 import json
 import time
 import torch
+
+# Check if CUDA is available and set the device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("Using device:", device)
+
 import argparse
 from PIL import Image
 import numpy as np
@@ -120,7 +125,7 @@ def main():
     model.eval()
 
     # Load Trained Weight #
-    device = torch.device("cuda") #vae.device()
+    
     if args.model.endswith(".pt") or args.model.endswith(".bin"):
         model.load_state_dict(torch.load(args.model), strict=False)
     else: