Selaa lähdekoodia

ONNX export explicit cpu map_location

5.0
Glenn Jocher 4 vuotta sitten
vanhempi
commit
1e2cb6b480
1 muutettua tiedostoa jossa 1 lisäystä ja 1 poistoa
  1. +1
    -1
      models/onnx_export.py

+ 1
- 1
models/onnx_export.py Näytä tiedosto

@@ -25,7 +25,7 @@ if __name__ == '__main__':

# Load pytorch model
google_utils.attempt_download(opt.weights)
model = torch.load(opt.weights)['model']
model = torch.load(opt.weights, map_location=torch.device('cpu'))['model']
model.eval()
model.fuse()


Loading…
Peruuta
Tallenna