Train with branch paper

[linghu8812]

Hello! When I finish training with branch paper, it encountered with an error:

  File "train.py", line 562, in <module>
    train(hyp, opt, device, tb_writer, wandb)
  File "train.py", line 369, in train
    log_imgs=opt.log_imgs if wandb else 0)
  File "/home/linghu8812/yolor_ele/test.py", line 211, in test
    plot_images(img, output_to_target(output, width, height), paths, f, names)  # predictions
  File "/home/linghu8812/yolor_ele/utils/plots.py", line 108, in output_to_target
    return np.array(targets)
  File "/media/D/anaconda3/envs/deeplearning/lib/python3.7/site-packages/torch/tensor.py", line 630, in __array__
    return self.numpy()
TypeError: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

perhaps it should be this?

 return np.array(targets.cpu())

[WongKinYiu]

Yes, it can solve the problem. https://github.com/WongKinYiu/yolor/issues/16#issuecomment-853717447

[danielcrane]

I noticed the same thing - would it help if I sent a pull request to the paper branch to fix this issue @WongKinYiu ?

The problem seems to stem from here:

    if isinstance(output, torch.Tensor):
        output = output.cpu().numpy()

from what I've seen output is (at least sometimes) actually a ragged list of tensors of different shapes, so this line doesn't trigger.

An inelegant solution could be something like so:

   for i, o in enumerate(output):
       if o is not None:
            if isinstance(o, torch.Tensor):
                o = o.cpu().numpy()