codingutf8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch videomae model import copy import unittest import numpy as np from huggingfacehub import hfhubdownload from transformers import videomaeconfig from transformers models auto import getvalues from transformers testingutils import requiretorch requirevision slow torchdevice from transformers utils import cachedproperty istorchavailable isvisionavailable from testconfigurationcommon import configtester from testmodelingcommon import modeltestermixin floatstensor idstensor from testpipelinemixin import pipelinetestermixin if istorchavailable import torch from torch import nn from transformers import modelforvideoclassificationmapping videomaeforpretraining videomaeforvideoclassification videomaemodel from transformers models videomae modelingvideomae import videomaepretrainedmodelarchivelist if isvisionavailable from transformers import videomaeimageprocessor class videomaemodeltester def init self parent batchsize13 imagesize10 numchannels3 patchsize2 tubeletsize2 numframes2 istrainingtrue uselabelstrue hiddensize32 numhiddenlayers2 numattentionheads4 intermediatesize37 hiddenactgelu hiddendropoutprob0 1 attentionprobsdropoutprob0 1 typesequencelabelsize10 initializerrange0 02 maskratio0 9 scopenone self parent parent self batchsize batchsize self imagesize imagesize self numchannels numchannels self patchsize patchsize self tubeletsize tubeletsize self numframes numframes self istraining istraining self uselabels uselabels self hiddensize hiddensize self numhiddenlayers numhiddenlayers self numattentionheads numattentionheads self intermediatesize intermediatesize self hiddenact hiddenact self hiddendropoutprob hiddendropoutprob self attentionprobsdropoutprob attentionprobsdropoutprob self typesequencelabelsize typesequencelabelsize self initializerrange initializerrange self maskratio maskratio self scope scope in videomae the number of tokens equals numframestubeletsize numpatches per frame self numpatchesperframe imagesize patchsize 2 self seqlength numframes tubeletsize self numpatchesperframe use this variable to define boolmaskedpos self nummasks intmaskratio self seqlength def prepareconfigandinputsself pixelvalues floatstensor self batchsize self numframes self numchannels self imagesize self imagesize labels none if self uselabels labels idstensorself batchsize self typesequencelabelsize config self getconfig return config pixelvalues labels def getconfigself return videomaeconfig imagesizeself imagesize patchsizeself patchsize numchannelsself numchannels numframesself numframes tubeletsizeself tubeletsize hiddensizeself hiddensize numhiddenlayersself numhiddenlayers numattentionheadsself numattentionheads intermediatesizeself intermediatesize hiddenactself hiddenact hiddendropoutprobself hiddendropoutprob attentionprobsdropoutprobself attentionprobsdropoutprob isdecoderfalse initializerrangeself initializerrange decoderhiddensizeself hiddensize decoderintermediatesizeself intermediatesize decodernumattentionheadsself numattentionheads decodernumhiddenlayersself numhiddenlayers def createandcheckmodelself config pixelvalues labels model videomaemodelconfigconfig model totorchdevice model eval result modelpixelvalues self parent assertequalresult lasthiddenstate shape self batchsize self seqlength self hiddensize def createandcheckforpretrainingself config pixelvalues labels model videomaeforpretrainingconfig model totorchdevice model eval important each video needs to have the same number of masked patches hence we define a single mask which we then repeat for each example in the batch mask torch onesself nummasks mask torch catmask torch zerosself seqlength mask size0 boolmaskedpos mask expandself batchsize 1 bool result modelpixelvalues boolmaskedpos model only returns predictions for masked patches nummaskedpatches mask sum item decodernumlabels 3 self tubeletsize self patchsize2 self parent assertequalresult logits shape self batchsize nummaskedpatches decodernumlabels def prepareconfigandinputsforcommonself configandinputs self prepareconfigandinputs config pixelvalues labels configandinputs inputsdict pixelvalues pixelvalues return config inputsdict requiretorch class videomaemodeltestmodeltestermixin pipelinetestermixin unittest testcase allmodelclasses videomaemodel videomaeforpretraining videomaeforvideoclassification if istorchavailable else pipelinemodelmapping featureextraction videomaemodel videoclassification videomaeforvideoclassification if istorchavailable else testpruning false testtorchscript false testresizeembeddings false testheadmasking false def setupself self modeltester videomaemodeltesterself self configtester configtesterself configclassvideomaeconfig hastextmodalityfalse hiddensize37 def prepareforclassself inputsdict modelclass returnlabelsfalse inputsdict copy deepcopyinputsdict if modelclass videomaeforpretraining important each video needs to have the same number of masked patches hence we define a single mask which we then repeat for each example in the batch mask torch onesself modeltester nummasks mask torch catmask torch zerosself modeltester seqlength mask size0 boolmaskedpos mask expandself modeltester batchsize 1 bool inputsdictboolmaskedpos boolmaskedpos totorchdevice if returnlabels if modelclass in getvaluesmodelforvideoclassificationmapping inputsdictlabels torch zeros self modeltester batchsize dtypetorch long devicetorchdevice return inputsdict def testconfigself self configtester runcommontests unittest skipreasonvideomae does not use inputsembeds def testinputsembedsself pass def testmodelcommonattributesself config self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses model modelclassconfig self assertisinstancemodel getinputembeddings nn module x model getoutputembeddings self asserttruex is none or isinstancex nn linear def testmodelself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckmodelconfigandinputs def testforpretrainingself configandinputs self modeltester prepareconfigandinputs self modeltester createandcheckforpretrainingconfigandinputs slow def testmodelfrompretrainedself for modelname in videomaepretrainedmodelarchivelist 1 model videomaemodel frompretrainedmodelname self assertisnotnonemodel def testattentionoutputsself if not self hasattentions pass else config inputsdict self modeltester prepareconfigandinputsforcommon config returndict true for modelclass in self allmodelclasses numvisiblepatches self modeltester seqlength self modeltester nummasks seqlen numvisiblepatches if modelclass videomaeforpretraining else self modeltester seqlength inputsdictoutputattentions true inputsdictoutputhiddenstates false config returndict true model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass attentions outputs attentions self assertequallenattentions self modeltester numhiddenlayers check that outputattentions also work using config del inputsdictoutputattentions config outputattentions true model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass attentions outputs attentions self assertequallenattentions self modeltester numhiddenlayers self assertlistequal listattentions0 shape3 self modeltester numattentionheads seqlen seqlen outlen lenoutputs check attention is always last and order is fine inputsdictoutputattentions true inputsdictoutputhiddenstates true model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass self assertequaloutlen 1 lenoutputs selfattentions outputs attentions self assertequallenselfattentions self modeltester numhiddenlayers self assertlistequal listselfattentions0 shape3 self modeltester numattentionheads seqlen seqlen def testhiddenstatesoutputself def checkhiddenstatesoutputinputsdict config modelclass model modelclassconfig model totorchdevice model eval with torch nograd outputs modelself prepareforclassinputsdict modelclass hiddenstates outputs hiddenstates expectednumlayers self modeltester numhiddenlayers 1 self assertequallenhiddenstates expectednumlayers numvisiblepatches self modeltester seqlength self modeltester nummasks seqlength numvisiblepatches if modelclass videomaeforpretraining else self modeltester seqlength self assertlistequal listhiddenstates0 shape2 seqlength self modeltester hiddensize config inputsdict self modeltester prepareconfigandinputsforcommon for modelclass in self allmodelclasses inputsdictoutputhiddenstates true checkhiddenstatesoutputinputsdict config modelclass check that outputhiddenstates also work using config del inputsdictoutputhiddenstates config outputhiddenstates true checkhiddenstatesoutputinputsdict config modelclass we will verify our results on a video of eating spaghetti frame indices used 164 168 172 176 181 185 189 193 198 202 206 210 215 219 223 227 def preparevideo file hfhubdownload repoidhfinternaltestingspaghettivideo filenameeatingspaghetti npy repotypedataset video np loadfile return listvideo requiretorch requirevision class videomaemodelintegrationtestunittest testcase cachedproperty def defaultimageprocessorself logits were tested with a different mean and std so we use the same here return videomaeimageprocessorimagemean0 5 0 5 0 5 imagestd0 5 0 5 0 5 if isvisionavailable else none slow def testinferenceforvideoclassificationself model videomaeforvideoclassification frompretrainedmcgnjuvideomaebasefinetunedkinetics to torchdevice imageprocessor self defaultimageprocessor video preparevideo inputs imageprocessorvideo returntensorspt totorchdevice forward pass with torch nograd outputs modelinputs verify the logits expectedshape torch size1 400 self assertequaloutputs logits shape expectedshape expectedslice torch tensor0 3669 0 0688 0 2421 totorchdevice self asserttruetorch allcloseoutputs logits0 3 expectedslice atol1e4 slow def testinferenceforpretrainingself model videomaeforpretraining frompretrainedmcgnjuvideomaebaseshort totorchdevice imageprocessor self defaultimageprocessor video preparevideo inputs imageprocessorvideo returntensorspt totorchdevice add boolean mask indicating which patches to mask localpath hfhubdownloadrepoidhfinternaltestingboolmaskedpos filenameboolmaskedpos pt inputsboolmaskedpos torch loadlocalpath forward pass with torch nograd outputs modelinputs verify the logits expectedshape torch size1 1408 1536 expectedslice torch tensor 0 7994 0 9612 0 8508 0 7401 0 8958 0 8302 0 5862 0 7468 0 7325 devicetorchdevice self assertequaloutputs logits shape expectedshape self asserttruetorch allcloseoutputs logits0 3 3 expectedslice atol1e4 verify the loss config normpixloss true expectedloss torch tensor0 5142 devicetorchdevice self asserttruetorch allcloseoutputs loss expectedloss atol1e4 verify the loss config normpixloss false model videomaeforpretraining frompretrainedmcgnjuvideomaebaseshort normpixlossfalse to torchdevice with torch nograd outputs modelinputs expectedloss torch tensortorch tensor0 6469 devicetorchdevice self asserttruetorch allcloseoutputs loss expectedloss atol1e4 coding utf 8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license testing suite for the pytorch videomae model in videomae the number of tokens equals num_frames tubelet_size num_patches per frame use this variable to define bool_masked_pos important each video needs to have the same number of masked patches hence we define a single mask which we then repeat for each example in the batch model only returns predictions for masked patches here we also overwrite some of the tests of test_modeling_common py as videomae does not use input_ids inputs_embeds attention_mask and seq_length important each video needs to have the same number of masked patches hence we define a single mask which we then repeat for each example in the batch check that output_attentions also work using config check attention is always last and order is fine check that output_hidden_states also work using config we will verify our results on a video of eating spaghetti frame indices used 164 168 172 176 181 185 189 193 198 202 206 210 215 219 223 227 logits were tested with a different mean and std so we use the same here forward pass verify the logits add boolean mask indicating which patches to mask forward pass verify the logits verify the loss config norm_pix_loss true verify the loss config norm_pix_loss false
import copy
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import VideoMAEConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch
    from torch import nn

    from transformers import (
        MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING,
        VideoMAEForPreTraining,
        VideoMAEForVideoClassification,
        VideoMAEModel,
    )
    from transformers.models.videomae.modeling_videomae import VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from transformers import VideoMAEImageProcessor


class VideoMAEModelTester:
    def __init__(
        self,
        parent,
        batch_size=13,
        image_size=10,
        num_channels=3,
        patch_size=2,
        tubelet_size=2,
        num_frames=2,
        is_training=True,
        use_labels=True,
        hidden_size=32,
        num_hidden_layers=2,
        num_attention_heads=4,
        intermediate_size=37,
        hidden_act="gelu",
        hidden_dropout_prob=0.1,
        attention_probs_dropout_prob=0.1,
        type_sequence_label_size=10,
        initializer_range=0.02,
        mask_ratio=0.9,
        scope=None,
    ):
        self.parent = parent
        self.batch_size = batch_size
        self.image_size = image_size
        self.num_channels = num_channels
        self.patch_size = patch_size
        self.tubelet_size = tubelet_size
        self.num_frames = num_frames
        self.is_training = is_training
        self.use_labels = use_labels
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.hidden_act = hidden_act
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.type_sequence_label_size = type_sequence_label_size
        self.initializer_range = initializer_range
        self.mask_ratio = mask_ratio
        self.scope = scope

        
        self.num_patches_per_frame = (image_size // patch_size) ** 2
        self.seq_length = (num_frames // tubelet_size) * self.num_patches_per_frame

        
        self.num_masks = int(mask_ratio * self.seq_length)

    def prepare_config_and_inputs(self):
        pixel_values = floats_tensor(
            [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]
        )

        labels = None
        if self.use_labels:
            labels = ids_tensor([self.batch_size], self.type_sequence_label_size)

        config = self.get_config()

        return config, pixel_values, labels

    def get_config(self):
        return VideoMAEConfig(
            image_size=self.image_size,
            patch_size=self.patch_size,
            num_channels=self.num_channels,
            num_frames=self.num_frames,
            tubelet_size=self.tubelet_size,
            hidden_size=self.hidden_size,
            num_hidden_layers=self.num_hidden_layers,
            num_attention_heads=self.num_attention_heads,
            intermediate_size=self.intermediate_size,
            hidden_act=self.hidden_act,
            hidden_dropout_prob=self.hidden_dropout_prob,
            attention_probs_dropout_prob=self.attention_probs_dropout_prob,
            is_decoder=False,
            initializer_range=self.initializer_range,
            decoder_hidden_size=self.hidden_size,
            decoder_intermediate_size=self.intermediate_size,
            decoder_num_attention_heads=self.num_attention_heads,
            decoder_num_hidden_layers=self.num_hidden_layers,
        )

    def create_and_check_model(self, config, pixel_values, labels):
        model = VideoMAEModel(config=config)
        model.to(torch_device)
        model.eval()
        result = model(pixel_values)
        self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))

    def create_and_check_for_pretraining(self, config, pixel_values, labels):
        model = VideoMAEForPreTraining(config)
        model.to(torch_device)
        model.eval()
        
        
        mask = torch.ones((self.num_masks,))
        mask = torch.cat([mask, torch.zeros(self.seq_length - mask.size(0))])
        bool_masked_pos = mask.expand(self.batch_size, -1).bool()

        result = model(pixel_values, bool_masked_pos)
        
        num_masked_patches = mask.sum().item()
        decoder_num_labels = 3 * self.tubelet_size * self.patch_size**2
        self.parent.assertEqual(result.logits.shape, (self.batch_size, num_masked_patches, decoder_num_labels))

    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        config, pixel_values, labels = config_and_inputs
        inputs_dict = {"pixel_values": pixel_values}
        return config, inputs_dict


@require_torch
class VideoMAEModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
    

    all_model_classes = (
        (VideoMAEModel, VideoMAEForPreTraining, VideoMAEForVideoClassification) if is_torch_available() else ()
    )
    pipeline_model_mapping = (
        {"feature-extraction": VideoMAEModel, "video-classification": VideoMAEForVideoClassification}
        if is_torch_available()
        else {}
    )

    test_pruning = False
    test_torchscript = False
    test_resize_embeddings = False
    test_head_masking = False

    def setUp(self):
        self.model_tester = VideoMAEModelTester(self)
        self.config_tester = ConfigTester(self, config_class=VideoMAEConfig, has_text_modality=False, hidden_size=37)

    def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
        inputs_dict = copy.deepcopy(inputs_dict)

        if model_class == VideoMAEForPreTraining:
            
            
            mask = torch.ones((self.model_tester.num_masks,))
            mask = torch.cat([mask, torch.zeros(self.model_tester.seq_length - mask.size(0))])
            bool_masked_pos = mask.expand(self.model_tester.batch_size, -1).bool()
            inputs_dict["bool_masked_pos"] = bool_masked_pos.to(torch_device)

        if return_labels:
            if model_class in [
                *get_values(MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING),
            ]:
                inputs_dict["labels"] = torch.zeros(
                    self.model_tester.batch_size, dtype=torch.long, device=torch_device
                )

        return inputs_dict

    def test_config(self):
        self.config_tester.run_common_tests()

    @unittest.skip(reason="VideoMAE does not use inputs_embeds")
    def test_inputs_embeds(self):
        pass

    def test_model_common_attributes(self):
        config, _ = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model = model_class(config)
            self.assertIsInstance(model.get_input_embeddings(), (nn.Module))
            x = model.get_output_embeddings()
            self.assertTrue(x is None or isinstance(x, nn.Linear))

    def test_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*config_and_inputs)

    def test_for_pretraining(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_pretraining(*config_and_inputs)

    @slow
    def test_model_from_pretrained(self):
        for model_name in VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            model = VideoMAEModel.from_pretrained(model_name)
            self.assertIsNotNone(model)

    def test_attention_outputs(self):
        if not self.has_attentions:
            pass

        else:
            config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
            config.return_dict = True

            for model_class in self.all_model_classes:
                num_visible_patches = self.model_tester.seq_length - self.model_tester.num_masks
                seq_len = (
                    num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length
                )

                inputs_dict["output_attentions"] = True
                inputs_dict["output_hidden_states"] = False
                config.return_dict = True
                model = model_class(config)
                model.to(torch_device)
                model.eval()
                with torch.no_grad():
                    outputs = model(**self._prepare_for_class(inputs_dict, model_class))
                attentions = outputs.attentions
                self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)

                
                del inputs_dict["output_attentions"]
                config.output_attentions = True
                model = model_class(config)
                model.to(torch_device)
                model.eval()
                with torch.no_grad():
                    outputs = model(**self._prepare_for_class(inputs_dict, model_class))
                attentions = outputs.attentions
                self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)

                self.assertListEqual(
                    list(attentions[0].shape[-3:]),
                    [self.model_tester.num_attention_heads, seq_len, seq_len],
                )
                out_len = len(outputs)

                
                inputs_dict["output_attentions"] = True
                inputs_dict["output_hidden_states"] = True
                model = model_class(config)
                model.to(torch_device)
                model.eval()
                with torch.no_grad():
                    outputs = model(**self._prepare_for_class(inputs_dict, model_class))

                self.assertEqual(out_len + 1, len(outputs))

                self_attentions = outputs.attentions

                self.assertEqual(len(self_attentions), self.model_tester.num_hidden_layers)
                self.assertListEqual(
                    list(self_attentions[0].shape[-3:]),
                    [self.model_tester.num_attention_heads, seq_len, seq_len],
                )

    def test_hidden_states_output(self):
        def check_hidden_states_output(inputs_dict, config, model_class):
            model = model_class(config)
            model.to(torch_device)
            model.eval()

            with torch.no_grad():
                outputs = model(**self._prepare_for_class(inputs_dict, model_class))

            hidden_states = outputs.hidden_states
            expected_num_layers = self.model_tester.num_hidden_layers + 1
            self.assertEqual(len(hidden_states), expected_num_layers)

            num_visible_patches = self.model_tester.seq_length - self.model_tester.num_masks
            seq_length = num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length

            self.assertListEqual(
                list(hidden_states[0].shape[-2:]),
                [seq_length, self.model_tester.hidden_size],
            )

        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            inputs_dict["output_hidden_states"] = True
            check_hidden_states_output(inputs_dict, config, model_class)

            
            del inputs_dict["output_hidden_states"]
            config.output_hidden_states = True

            check_hidden_states_output(inputs_dict, config, model_class)




def prepare_video():
    file = hf_hub_download(
        repo_id="hf-internal-testing/spaghetti-video", filename="eating_spaghetti.npy", repo_type="dataset"
    )
    video = np.load(file)
    return list(video)


@require_torch
@require_vision
class VideoMAEModelIntegrationTest(unittest.TestCase):
    @cached_property
    def default_image_processor(self):
        
        return (
            VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5], image_std=[0.5, 0.5, 0.5])
            if is_vision_available()
            else None
        )

    @slow
    def test_inference_for_video_classification(self):
        model = VideoMAEForVideoClassification.from_pretrained("MCG-NJU/videomae-base-finetuned-kinetics").to(
            torch_device
        )

        image_processor = self.default_image_processor
        video = prepare_video()
        inputs = image_processor(video, return_tensors="pt").to(torch_device)

        
        with torch.no_grad():
            outputs = model(**inputs)

        
        expected_shape = torch.Size((1, 400))
        self.assertEqual(outputs.logits.shape, expected_shape)

        expected_slice = torch.tensor([0.3669, -0.0688, -0.2421]).to(torch_device)

        self.assertTrue(torch.allclose(outputs.logits[0, :3], expected_slice, atol=1e-4))

    @slow
    def test_inference_for_pretraining(self):
        model = VideoMAEForPreTraining.from_pretrained("MCG-NJU/videomae-base-short").to(torch_device)

        image_processor = self.default_image_processor
        video = prepare_video()
        inputs = image_processor(video, return_tensors="pt").to(torch_device)

        
        local_path = hf_hub_download(repo_id="hf-internal-testing/bool-masked-pos", filename="bool_masked_pos.pt")
        inputs["bool_masked_pos"] = torch.load(local_path)

        
        with torch.no_grad():
            outputs = model(**inputs)

        
        expected_shape = torch.Size([1, 1408, 1536])
        expected_slice = torch.tensor(
            [[0.7994, 0.9612, 0.8508], [0.7401, 0.8958, 0.8302], [0.5862, 0.7468, 0.7325]], device=torch_device
        )
        self.assertEqual(outputs.logits.shape, expected_shape)
        self.assertTrue(torch.allclose(outputs.logits[0, :3, :3], expected_slice, atol=1e-4))

        
        expected_loss = torch.tensor([0.5142], device=torch_device)
        self.assertTrue(torch.allclose(outputs.loss, expected_loss, atol=1e-4))

        
        model = VideoMAEForPreTraining.from_pretrained("MCG-NJU/videomae-base-short", norm_pix_loss=False).to(
            torch_device
        )

        with torch.no_grad():
            outputs = model(**inputs)

        expected_loss = torch.tensor(torch.tensor([0.6469]), device=torch_device)
        self.assertTrue(torch.allclose(outputs.loss, expected_loss, atol=1e-4))