How to check device of torch.nn.Module?

I received the following error:

File ".../lib/python3.6/site-packages/torch/nn/", line 1852, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: Input, output and indices must be on the current device

So I wanted to check what devices the three variables were on. For the tensors, I could use tensor.get_device() and that worked fine. However, when I tried checking what device the problematic torch.nn.Module was on, I received the error:

torch.nn.modules.module.ModuleAttributeError: 'DistilBertForMaskedLM' object has no attribute 'get_device'

How does one check the device of a torch.nn.Module object?

I just found this Stack Overflow post. I understand now.

1 Like