Add `_no_split_modules`
Browse files- modeling_aimv2.py +1 -0
modeling_aimv2.py
CHANGED
@@ -149,6 +149,7 @@ class AIMv2PretrainedModel(PreTrainedModel):
|
|
149 |
config_class = AIMv2Config
|
150 |
base_model_prefix = "aimv2"
|
151 |
main_input_name = "pixel_values"
|
|
|
152 |
_supports_sdpa = True
|
153 |
|
154 |
|
|
|
149 |
config_class = AIMv2Config
|
150 |
base_model_prefix = "aimv2"
|
151 |
main_input_name = "pixel_values"
|
152 |
+
_no_split_modules = ["AIMv2ViTPreprocessor", "AIMv2Block"]
|
153 |
_supports_sdpa = True
|
154 |
|
155 |
|