UPDATE

app.py

@@ -34,13 +34,12 @@ CONFIG = Config()
 import numpy as np
 import torch
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-if device == "cpu" and CONFIG.STRATEGY != "cpu":
-    logger.info(f"Cuda not found, fall back to cpu")
+if "cuda" in CONFIG.STRATEGY.lower() and not torch.cuda.is_available():
+    logger.info(f"CUDA not found, fall back to cpu")
     CONFIG.STRATEGY = "cpu"
 
-if "cuda" in CONFIG.STRATEGY:
+if "cuda" in CONFIG.STRATEGY.lower():
     from pynvml import *
 
     nvmlInit()