Update modeling_flow_match.py

modeling_flow_match.py

@@ -337,10 +337,20 @@ class FlowMatchRelayModel(PreTrainedModel):
         images = model.sample(n_samples=8, class_label=3)
     """
     config_class = FlowMatchRelayConfig
+    _tied_weights_keys = []
+    _keys_to_ignore_on_load_missing = []
+    _keys_to_ignore_on_load_unexpected = []
+    _no_split_modules = []
+    supports_gradient_checkpointing = False
 
     def __init__(self, config):
         super().__init__(config)
         self.unet = FlowMatchUNet(config)
+        self.post_init()
+
+    def _init_weights(self, module):
+        """No-op — weights loaded from checkpoint or already initialized."""
+        pass
 
     def forward(self, x, t, class_labels):
         """