Unity-Technologies / Robotics-Object-Pose-Estimation

A complete end-to-end demonstration in which we collect training data in Unity and use that data to train a deep neural network to predict the pose of a cube. This model is then deployed in a simulated robotic pick-and-place task.
Apache License 2.0
293 stars 75 forks source link

AIRO-1654 fixed cpu error #50

Closed JonathanLeban closed 2 years ago

JonathanLeban commented 2 years ago

Proposed change(s)

Describe the changes made in this PR.

Fixed a cpu/gpu error made during the inference.

at669 commented 2 years ago

Linked to #45.

adakoda commented 2 years ago

Thank you for your work.

I had tested your code changes. But my cuda 11.6 gou environment could not solve the problem without adding following patch.

add caling model.to(device) after calling model.load_state_dict()

at669 commented 2 years ago

Thank you for your work.

I had tested your code changes. But my cuda 11.6 gou environment could not solve the problem without adding following patch.

add caling model.to(device) after calling model.load_state_dict()

Thank you for the feedback @adakoda. Could you please pull down #52 and rerun to see if this resolves your issue? I was able to reproduce the issue as stated before, and this new PR resolves it for me.