Deprecation warnings

3dem · Oct 24, 2024 · 457b82c · 457b82c
1 parent 8b3a686
commit 457b82c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/model_angelo/models/multi_gpu_wrapper.py b/model_angelo/models/multi_gpu_wrapper.py
@@ -93,7 +93,7 @@ def run_inference(
                 inference_data = input_queue.get()
                 if inference_data.status != 1:
                     break
-                with torch.cuda.amp.autocast(dtype=dtype):
+                with torch.autocast(device_type="cuda", dtype=dtype):
                     output = model(**inference_data.data)
                 output = output.to("cpu").to(torch.float32)
                 output_queue.put(output)
@@ -155,7 +155,7 @@ def forward(self, data_list: List) -> List:
                     InferenceData(data=send_dict_to_device(data, device), status=1)
                 )
             else:
-                with torch.cuda.amp.autocast(dtype=self.dtype), torch.no_grad():
+                with torch.autocast(device_type="cuda", dtype=self.dtype), torch.no_grad():
                     output_list.append(
                         self.model(**send_dict_to_device(data, device)).to("cpu").to(torch.float32)
                     )