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".

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Multi-GPU support

engrean opened this issue · comments

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

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