CRF layer fails when running on multiple GPUs

[erastogi]

I am running the following code:

 device = logits.device
 log_likelihood = self.crf(logits.to(device), crf_labels.to(device), mask=attention_mask.to(device).type(torch.ByteTensor))

However, I am getting the following error:

 File "/opt/conda/lib/python3.8/site-packages/torchcrf/__init__.py", line 102, in forward
 numerator = self._compute_score(emissions, tags, mask)
 File "/opt/conda/lib/python3.8/site-packages/torchcrf/__init__.py", line 187, in _compute_score
 score += emissions[0, torch.arange(batch_size), tags[0]]
 RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:4 and cpu!
 RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:5 and cpu!
 RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:3 and cpu!
 RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:6 and cpu!
 RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
 RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:2 and cpu!
 RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:7 and cpu!"

What would be the best way to run crf on multiple GPUs?

[kmkurn]

I have never worked with multiple GPUs so not sure if I can help on that, but from a glance at your code, you may want to change .type(torch.ByteTensor) to .bool() because torch.ByteTensor is for tensors on CPU. Let me know if this solves the issue.

[erastogi]

It worked, thanks a ton!!