ynhe commited on
Commit
6afc77b
β€’
1 Parent(s): ad1552b

Update models/videochat.py

Browse files
Files changed (1) hide show
  1. models/videochat.py +2 -0
models/videochat.py CHANGED
@@ -137,6 +137,8 @@ class VideoChat(Blip2Base):
137
  llama_model_path,
138
  torch_dtype=torch.float16,
139
  use_auth_token=os.environ["HF_TOKEN"],
 
 
140
  )
141
 
142
  print("freeze LLAMA")
 
137
  llama_model_path,
138
  torch_dtype=torch.float16,
139
  use_auth_token=os.environ["HF_TOKEN"],
140
+ load_in_8bit=True,
141
+ device_map="auto"
142
  )
143
 
144
  print("freeze LLAMA")