update

alibaba · Apr 9, 2024 · 9339150 · 9339150
1 parent d5ccb49
commit 9339150
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/easycv/apis/export.py b/easycv/apis/export.py
@@ -239,10 +239,11 @@ def _export_cls(model, cfg, filename):
     if model_config['backbone'].get(
             'type', None) == 'ResNet' and model_config['backbone'].get(
                 'depth', None) == 50:
-
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
         model.eval()
+        model.to(device)
         img_size = int(cfg.image_size2)
-        x_input = torch.randn((1, 3, img_size, img_size))
+        x_input = torch.randn((1, 3, img_size, img_size)).to(device)
         torch.onnx.export(
             model,
             (x_input, 'test'),