can't use pretrained weights when working with a tpu
Author: level14taken
Pull request: no

### System Info

transformers-4.32
python-3.8


### Who can help?

@Rocketknight1 @sanchit-gandhi@gante

### Information

- [x] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

1. The error happens only on kaggle tpu, and it works fine on colab
2. But, it just follows the tpu tutorial and is very basic.

```python
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu="local")
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)

from transformers import TFWav2Vec2Model,AutoModel
with strategy.scope():
    model = TFWav2Vec2Model.from_pretrained(model_checkpoint,from_pt=True) #####error
    model.compile(optimizer="adam")
```

### Expected behavior

Full stack trace of the error.
```python

TFWav2Vec2Model has backpropagation operations that are NOT supported on CPU. If you wish to train/fine-tune this model, you need a GPU or a TPU

---------------------------------------------------------------------------
NotFoundError                             Traceback (most recent call last)
Cell In[7], line 3
      1 from transformers import TFWav2Vec2Model,AutoModel
      2 with strategy.scope():
----> 3     model = TFWav2Vec2Model.from_pretrained(model_checkpoint,from_pt=True)
      4     # You can compile with jit_compile=True when debugging on CPU or GPU to check
      5     # that XLA compilation works. Remember to take it out when actually running
      6     # on TPU, though - XLA compilation will be handled for you when running with a
      7     # TPUStrategy!
      8     model.compile(optimizer="adamw")

File /usr/local/lib/python3.8/site-packages/transformers/modeling_tf_utils.py:2894, in TFPreTrainedModel.from_pretrained(cls, pretrained_model_name_or_path, config, cache_dir, ignore_mismatched_sizes, force_download, local_files_only, token, revision, *model_args, **kwargs)
   2891     from .modeling_tf_pytorch_utils import load_pytorch_checkpoint_in_tf2_model
   2893     # Load from a PyTorch checkpoint
-> 2894     return load_pytorch_checkpoint_in_tf2_model(
   2895         model,
   2896         resolved_archive_file,
   2897         allow_missing_keys=True,
   2898         output_loading_info=output_loading_info,
   2899         _prefix=load_weight_prefix,
   2900         tf_to_pt_weight_rename=tf_to_pt_weight_rename,
   2901     )
   2903 # we might need to extend the variable scope for composite models
   2904 if load_weight_prefix is not None:

File /usr/local/lib/python3.8/site-packages/transformers/modeling_tf_pytorch_utils.py:189, in load_pytorch_checkpoint_in_tf2_model(tf_model, pytorch_checkpoint_path, tf_inputs, allow_missing_keys, output_loading_info, _prefix, tf_to_pt_weight_rename)
    185     pt_state_dict.update(torch.load(pt_path, map_location="cpu"))
    187 logger.info(f"PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()):,} parameters")
--> 189 return load_pytorch_weights_in_tf2_model(
    190     tf_model,
    191     pt_state_dict,
    192     tf_inputs=tf_inputs,
    193     allow_missing_keys=allow_missing_keys,
    194     output_loading_info=output_loading_info,
    195     _prefix=_prefix,
    196     tf_to_pt_weight_rename=tf_to_pt_weight_rename,
    197 )

File /usr/local/lib/python3.8/site-packages/transformers/modeling_tf_pytorch_utils.py:230, in load_pytorch_weights_in_tf2_model(tf_model, pt_state_dict, tf_inputs, allow_missing_keys, output_loading_info, _prefix, tf_to_pt_weight_rename)
    227     raise
    229 pt_state_dict = {k: v.numpy() for k, v in pt_state_dict.items()}
--> 230 return load_pytorch_state_dict_in_tf2_model(
    231     tf_model,
    232     pt_state_dict,
    233     tf_inputs=tf_inputs,
    234     allow_missing_keys=allow_missing_keys,
    235     output_loading_info=output_loading_info,
    236     _prefix=_prefix,
    237     tf_to_pt_weight_rename=tf_to_pt_weight_rename,
    238 )

File /usr/local/lib/python3.8/site-packages/transformers/modeling_tf_pytorch_utils.py:263, in load_pytorch_state_dict_in_tf2_model(tf_model, pt_state_dict, tf_inputs, allow_missing_keys, output_loading_info, _prefix, tf_to_pt_weight_rename, ignore_mismatched_sizes)
    261 if tf_inputs:
    262     with tf.name_scope(_prefix):
--> 263         tf_model(tf_inputs, training=False)  # Make sure model is built
    264 # Convert old format to new format if needed from a PyTorch state_dict
    265 tf_keys_to_pt_keys = {}

File /usr/local/lib/python3.8/site-packages/keras/utils/traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     67     filtered_tb = _process_traceback_frames(e.__traceback__)
     68     # To get the full stack trace, call:
     69     # `tf.debugging.disable_traceback_filtering()`
---> 70     raise e.with_traceback(filtered_tb) from None
     71 finally:
     72     del filtered_tb

File /usr/local/lib/python3.8/site-packages/transformers/modeling_tf_utils.py:426, in unpack_inputs.<locals>.run_call_with_unpacked_inputs(self, *args, **kwargs)
    423     config = self.config
    425 unpacked_inputs = input_processing(func, config, **fn_args_and_kwargs)
--> 426 return func(self, **unpacked_inputs)

File /usr/local/lib/python3.8/site-packages/transformers/models/wav2vec2/modeling_tf_wav2vec2.py:1407, in TFWav2Vec2Model.call(self, input_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, output_attentions, output_hidden_states, return_dict, training)
   1404 output_attentions = output_attentions if output_attentions else self.config.output_attentions
   1405 return_dict = return_dict if return_dict else self.config.return_dict
-> 1407 outputs = self.wav2vec2(
   1408     input_values=input_values,
   1409     attention_mask=attention_mask,
   1410     token_type_ids=token_type_ids,
   1411     position_ids=position_ids,
   1412     head_mask=head_mask,
   1413     inputs_embeds=inputs_embeds,
   1414     output_attentions=output_attentions,
   1415     output_hidden_states=output_hidden_states,
   1416     return_dict=return_dict,
   1417     training=training,
   1418 )
   1420 return outputs

File /usr/local/lib/python3.8/site-packages/transformers/modeling_tf_utils.py:426, in unpack_inputs.<locals>.run_call_with_unpacked_inputs(self, *args, **kwargs)
    423     config = self.config
    425 unpacked_inputs = input_processing(func, config, **fn_args_and_kwargs)
--> 426 return func(self, **unpacked_inputs)

File /usr/local/lib/python3.8/site-packages/transformers/models/wav2vec2/modeling_tf_wav2vec2.py:1158, in TFWav2Vec2MainLayer.call(self, input_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, output_attentions, output_hidden_states, return_dict, training, **kwargs)
   1155 if training:
   1156     hidden_states = self._mask_hidden_states(hidden_states, mask_time_indices=mask_time_indices)
-> 1158 encoder_outputs = self.encoder(
   1159     hidden_states,
   1160     attention_mask=attention_mask,
   1161     output_attentions=output_attentions,
   1162     output_hidden_states=output_hidden_states,
   1163     return_dict=return_dict,
   1164     training=training,
   1165 )
   1166 hidden_states = encoder_outputs[0]
   1168 if not return_dict:

File /usr/local/lib/python3.8/site-packages/transformers/models/wav2vec2/modeling_tf_wav2vec2.py:1007, in TFWav2Vec2EncoderStableLayerNorm.call(self, hidden_states, attention_mask, output_attentions, output_hidden_states, return_dict, training)
   1004 else:
   1005     attention_mask = None
-> 1007 position_embeddings = self.pos_conv_embed(hidden_states)
   1008 hidden_states = hidden_states + position_embeddings
   1009 hidden_states = self.dropout(hidden_states, training=training)

File /usr/local/lib/python3.8/site-packages/transformers/models/wav2vec2/modeling_tf_wav2vec2.py:568, in TFWav2Vec2PositionalConvEmbedding.call(self, hidden_states)
    567 def call(self, hidden_states: tf.Tensor) -> tf.Tensor:
--> 568     hidden_states = self.conv(hidden_states)
    569     hidden_states = self.padding(hidden_states)
    570     hidden_states = self.activation(hidden_states)

File /usr/local/lib/python3.8/site-packages/transformers/models/wav2vec2/modeling_tf_wav2vec2.py:480, in TFWav2Vec2WeightNormConv1D.call(self, inputs)
    477 self._normalize_kernel()
    479 padded_inputs = tf.pad(inputs, ((0, 0), (self.explicit_padding, self.explicit_padding), (0, 0)))
--> 480 output = super().call(padded_inputs)
    482 return output

NotFoundError: Exception encountered when calling layer 'conv' (type TFWav2Vec2WeightNormConv1D).

could not find registered transfer manager for platform Host -- check target linkage [Op:__inference__jit_compiled_convolution_op_1397]

Call arguments received by layer 'conv' (type TFWav2Vec2WeightNormConv1D):
  • inputs=tf.Tensor(shape=(1, 1, 1024), dtype=float32)

[ ]:

```

Comments:
Hm, if it works on Colab but not Kaggle, the issue is probably not in `transformers` itself. It's either in the versions of TF / transformers you're using on both platforms, or the way you're connecting to TPU on Kaggle. Maybe try `tpu = tf.distribute.cluster_resolver.TPUClusterResolver() ` instead of using `tpu='local'`?

Hi @Rocketknight1,
Thanks for the reply.

I've checked the versions(tranformers,tf) in colab and kaggle and they were exactly the same.
I've also checked with ``` tpu = tf.distribute.cluster_resolver.TPUClusterResolver() ``` but same issue came up. 
By the way kaggle uses TPU v3 and colab is using v2 
Does the error(registered transfer manager....) really mean something in huggingface or is it just pointing to a tensorflow op?

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.