Unverified Commit 1bbad7a2 authored by S.Kishore's avatar S.Kishore Committed by GitHub
Browse files

Added Type hints for VIT MAE (#19085)

* Added Type hints for VIT MAE

* Ran make fixup
parent fbe8464b
......@@ -665,13 +665,13 @@ class ViTMAEModel(ViTMAEPreTrainedModel):
@replace_return_docstrings(output_type=ViTMAEModelOutput, config_class=_CONFIG_FOR_DOC)
def forward(
self,
pixel_values=None,
noise=None,
head_mask=None,
output_attentions=None,
output_hidden_states=None,
return_dict=None,
):
pixel_values: Optional[torch.FloatTensor] = None,
noise: Optional[torch.FloatTensor] = None,
head_mask: Optional[torch.FloatTensor] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
) -> Union[Tuple, ViTMAEModelOutput]:
r"""
Returns:
......@@ -957,13 +957,13 @@ class ViTMAEForPreTraining(ViTMAEPreTrainedModel):
@replace_return_docstrings(output_type=ViTMAEForPreTrainingOutput, config_class=_CONFIG_FOR_DOC)
def forward(
self,
pixel_values=None,
noise=None,
head_mask=None,
output_attentions=None,
output_hidden_states=None,
return_dict=None,
):
pixel_values: Optional[torch.FloatTensor] = None,
noise: Optional[torch.FloatTensor] = None,
head_mask: Optional[torch.FloatTensor] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
) -> Union[Tuple, ViTMAEForPreTrainingOutput]:
r"""
Returns:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment