Update main_onnx_amd.py

This commit is contained in:
Elijah Harmon 2023-03-29 02:42:10 -04:00 committed by GitHub
parent 57513001ac
commit a56299f5e6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -116,7 +116,7 @@ def main():
if len(im.shape) == 3: if len(im.shape) == 3:
im = im[None] im = im[None]
outputs = ort_sess.run(None, {'images': npImg}) outputs = ort_sess.run(None, {'images': im})
im = torch.from_numpy(outputs[0]).to('cpu') im = torch.from_numpy(outputs[0]).to('cpu')