Revert changes

configuration_x3d.py

@@ -1,9 +1,15 @@
 from transformers import PretrainedConfig
-from x3d_model.cfg import load_config
+from .cfg import load_config
+
 
 class X3DConfig(PretrainedConfig):
     model_type = "x3d"
 
-    def __init__(self, path: str = None, **kwargs):
+    def __init__(self, **kwargs):
         super().__init__(**kwargs)
+
+        path = kwargs.get("path", None)
+        gpu_num = kwargs.get("gpu_num", 0)
+
         self.cfg = load_config(path)
+        self.cfg.NUM_GPUS = gpu_num

modeling_x3d.py

@@ -1,15 +1,23 @@
+import torch
 from transformers import PreTrainedModel
-from x3d_model.configuration_x3d import X3DConfig
-from x3d_model.x3d import build_model
+from .configuration_x3d import X3DConfig
+from .x3d import build_model
 
 
 class X3DModel(PreTrainedModel):
     config_class = X3DConfig
 
-    def __init__(self, config):
+    def __init__(self, config, **kwargs):
         super().__init__(config)
         self.model = build_model(config.cfg)
 
+        checkpoint = kwargs.get("checkpoint", None)
+
+        if checkpoint:
+            checkpoint = torch.load(
+                checkpoint, weights_only=True, map_location=torch.device("cpu"))
+            self.model.load_state_dict(checkpoint["model_state"])
+
     def forward(self, input_video):
         outputs = self.model(input_video)
         return outputs