Update vision_tower_builder.py

vision_tower_builder.py

@@ -2,9 +2,6 @@ from typing import Optional, Tuple, Union, Dict
 from dataclasses import dataclass
 from functools import partial, reduce
 from PIL import Image
-import torch
-import torch.utils.checkpoint
-from torch import nn
 import os
 from transformers.image_processing_utils import BatchFeature, get_size_dict
 from transformers.image_transforms import (
@@ -29,7 +26,7 @@ try:
     from flash_attn import flash_attn_qkvpacked_func
 except:
     print("You need to install flash_attn")
-from timm.models.layers import drop_path, to_2tuple, trunc_normal_
+from timm.layers import drop_path, to_2tuple, trunc_normal_
 
 
 
@@ -516,7 +513,7 @@ def build_vit(config, pt_type='origin'):
         drop_path_rate=0., 
         num_frames=config.num_frames,
         tubelet_size=1,
-        use_checkpoint=True,
+        use_checkpoint=False,
         checkpoint_num=24,
         return_index=config.return_idx,
         with_ln=True, # merge vision_layernorm in it
@@ -619,4 +616,4 @@ def build_vision_tower(vision_tower_cfg, **kwargs):
         raise NotImplementedError
         return UMTVisionTower(vision_tower, vision_tower_cfg=vision_tower_cfg, **kwargs)
 
-    raise ValueError(f"Unknown vision tower: {vision_tower}")
+    raise ValueError(f"Unknown vision tower: {vision_tower}")