| from torch import Tensor | |
| from transformers import PreTrainedModel | |
| from .configuration_act_estimator import ActEstimatorConfig | |
| from .model import VideoActionEstimator | |
| class ActEstimator(PreTrainedModel): | |
| config_class = ActEstimatorConfig | |
| def __init__(self, config: ActEstimatorConfig): | |
| super().__init__(config) | |
| self.model = VideoActionEstimator(**config.to_dict()) | |
| def forward(self, frames: Tensor, timestamps: Tensor = None) -> dict[str, Tensor]: | |
| return self.model(frames, timestamps) | |