In test.py just change this line
    plot_images(img, output_to_target(output, width, height), paths, str(f), names)  # predictions

to this
    plot_images(img, output_to_target(output.cpu(), width, height), paths, str(f), names)  # predictions