rsommerfeld / trocr

Powerful handwritten text recognition. A simple-to-use, unofficial implementation of the paper "TrOCR: Transformer-based Optical Character Recognition with Pre-trained Models".
MIT License
176 stars 27 forks source link

Multi-GPU support #5

Open engrean opened 2 years ago

engrean commented 2 years ago

I noticed that only one of the GPUs is being used when I train. I tried setting model to torch.nn.DataParallel(model), but I kept getting "RuntimeError: grad can be implicitly created only for scalar outputs". I am not quite familiar enough with Torch's multi-GPU support to fix it just yet. While I am looking, I am hoping someone could help make it support multiple GPUs

dparres commented 1 year ago

Hi, try to calculate loss and then do: loss = loss.sum() This will solve this problem.