[ViTHybrid] fix last accelerate slow test (#20705)

* fix last slow test

* revert deletion

* Update src/transformers/models/vit_hybrid/modeling_vit_hybrid.py
This commit is contained in:
Younes Belkada 2022-12-09 16:46:32 +01:00 committed by GitHub
parent 7319850902
commit 74330083b5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -358,7 +358,6 @@ class ViTHybridOutput(nn.Module):
return hidden_states
# Copied from transformers.models.vit.modeling_vit.ViTLayer with ViT->ViTHybrid
class ViTHybridLayer(nn.Module):
"""This corresponds to the Block class in the timm implementation."""
@ -387,7 +386,8 @@ class ViTHybridLayer(nn.Module):
outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
# first residual connection
hidden_states = attention_output + hidden_states
# We assign to correct device for `accelerate`, check: https://github.com/huggingface/transformers/pull/20705/
hidden_states = attention_output + hidden_states.to(attention_output.device)
# in ViTHybrid, layernorm is also applied after self-attention
layer_output = self.layernorm_after(hidden_states)