pytorch summary fails with huggingface model II: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu
I want a summary of a PyTorch model downloaded from huggingface:
from torchinfo import summary
from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)
summary(model, input_size=(16, 512), dtypes=['torch.IntTensor'])

(See SO for why the dtypes is needed.)
However, I am getting the error Expected all tensors to be on the same device, ... even though I have not provided any tensors. See the output below.
How can I fix this?


---------------------------------------------------------------------------

RuntimeError                              Traceback (most recent call last)

/usr/local/lib/python3.7/dist-packages/torchinfo/torchinfo.py in forward_pass(model, x, batch_dim, cache_forward_pass, device, **kwargs)
    257             if isinstance(x, (list, tuple)):
--&gt; 258                 _ = model.to(device)(*x, **kwargs)
    259             elif isinstance(x, dict):

11 frames

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-&gt; 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used

/usr/local/lib/python3.7/dist-packages/transformers/models/bert/modeling_bert.py in forward(self, input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, labels, output_attentions, output_hidden_states, return_dict)
   1530             output_hidden_states=output_hidden_states,
-&gt; 1531             return_dict=return_dict,
   1532         )

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1070 
-&gt; 1071         result = forward_call(*input, **kwargs)
   1072         if _global_forward_hooks or self._forward_hooks:

/usr/local/lib/python3.7/dist-packages/transformers/models/bert/modeling_bert.py in forward(self, input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, past_key_values, use_cache, output_attentions, output_hidden_states, return_dict)
    988             inputs_embeds=inputs_embeds,
--&gt; 989             past_key_values_length=past_key_values_length,
    990         )

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1070 
-&gt; 1071         result = forward_call(*input, **kwargs)
   1072         if _global_forward_hooks or self._forward_hooks:

/usr/local/lib/python3.7/dist-packages/transformers/models/bert/modeling_bert.py in forward(self, input_ids, token_type_ids, position_ids, inputs_embeds, past_key_values_length)
    214         if inputs_embeds is None:
--&gt; 215             inputs_embeds = self.word_embeddings(input_ids)
    216         token_type_embeddings = self.token_type_embeddings(token_type_ids)

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1070 
-&gt; 1071         result = forward_call(*input, **kwargs)
   1072         if _global_forward_hooks or self._forward_hooks:

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/sparse.py in forward(self, input)
    159             input, self.weight, self.padding_idx, self.max_norm,
--&gt; 160             self.norm_type, self.scale_grad_by_freq, self.sparse)
    161 

/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in embedding(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse)
   2042         _no_grad_embedding_renorm_(weight, input, max_norm, norm_type)
-&gt; 2043     return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
   2044 

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking arugment for argument index in method wrapper_index_select)


The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)

&lt;ipython-input-13-d6f4e53beef7&gt; in &lt;module&gt;()
      3 else:
      4     # Can't get this working. See https://stackoverflow.com/questions/68577198/pytorch-summary-fails-with-huggingface-model
----&gt; 5     summary(model, input_size=(16, 512), dtypes=['torch.IntTensor'])
      6     print(model)

/usr/local/lib/python3.7/dist-packages/torchinfo/torchinfo.py in summary(model, input_size, input_data, batch_dim, cache_forward_pass, col_names, col_width, depth, device, dtypes, row_settings, verbose, **kwargs)
    190     )
    191     summary_list = forward_pass(
--&gt; 192         model, x, batch_dim, cache_forward_pass, device, **kwargs
    193     )
    194     formatting = FormattingOptions(depth, verbose, col_names, col_width, row_settings)

/usr/local/lib/python3.7/dist-packages/torchinfo/torchinfo.py in forward_pass(model, x, batch_dim, cache_forward_pass, device, **kwargs)
    268             &quot;Failed to run torchinfo. See above stack traces for more details. &quot;
    269             f&quot;Executed layers up to: {executed_layers}&quot;
--&gt; 270         ) from e
    271     finally:
    272         if hooks is not None:

RuntimeError: Failed to run torchinfo. See above stack traces for more details. Executed layers up to: []


Output from transformers-cli:
- `transformers` version: 4.9.1
- Platform: Linux-5.4.104+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.11
- PyTorch version (GPU?): 1.9.0+cu102 (True)
- Tensorflow version (GPU?): 2.5.0 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: &lt;fill in&gt;
- Using distributed or parallel set-up in script?: &lt;fill in&gt;


A working solution (or workaround?) is kind of obvious:
summary(model, input_size=(16, 512), dtypes=['torch.IntTensor'], device='cpu')

