r3gm commited on
Commit
0422656
1 Parent(s): 3e3db38

Update src/mdx.py

Browse files
Files changed (1) hide show
  1. src/mdx.py +4 -3
src/mdx.py CHANGED
@@ -66,9 +66,10 @@ class MDX:
66
 
67
  # Set the device and the provider (CPU or CUDA)
68
  #self.device = torch.device(f'cuda:{processor}') if processor >= 0 else torch.device('cpu')
69
- device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
70
- self.provider = ['CUDAExecutionProvider'] if processor >= 0 else ['CPUExecutionProvider']
71
-
 
72
  self.model = params
73
 
74
  # Load the ONNX model using ONNX Runtime
 
66
 
67
  # Set the device and the provider (CPU or CUDA)
68
  #self.device = torch.device(f'cuda:{processor}') if processor >= 0 else torch.device('cpu')
69
+ self.device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
70
+ #self.provider = ['CUDAExecutionProvider'] if processor >= 0 else ['CPUExecutionProvider']
71
+ self.provider = ['CPUExecutionProvider']
72
+
73
  self.model = params
74
 
75
  # Load the ONNX model using ONNX Runtime