from transformers import PreTrainedModel
from .configuration_x3d import X3DConfig
from .x3d import build_model


class X3DModel(PreTrainedModel):
    config_class = X3DConfig

    def __init__(self, config):
        super().__init__(config)
        self.model = build_model(config.cfg)

    def forward(self, input_video):
        outputs = self.model(input_video)
        return outputs