michalk8 commited on
Commit
543f2f4
1 Parent(s): 7a8d3fc

Add `_no_split_modules`

Browse files
Files changed (1) hide show
  1. 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