demo.py freezes on model prediction

[MihailMihaylov97]

Running python3 demo.py --checkpoint=data/model_checkpoint.pt --img=examples/im1010.jpg --bbox=examples/im1010_bbox.json in a docker environment using the provided Docker image results in script freezing after preprocessing the image, at the model prediction step. The container keeps running, but the code is frozen. Eventually may load the model, but immediately afterwards fails with

RuntimeError: cublas runtime error : the GPU program failed to execute at /tmp/pip-req-build-5vv9db21/aten/src/THC/THCBlas.cu:450

Update: Same if I set device to 'cpu'.