chenkq commited on
Commit
2ecffe1
1 Parent(s): 8abca87

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +1 -1
README.md CHANGED
@@ -109,7 +109,7 @@ with init_empty_weights():
109
  low_cpu_mem_usage=True,
110
  trust_remote_code=True,
111
  )
112
- device_map = infer_auto_device_map(model, max_memory={0:'20GiB',1:'20GiB','cpu':'16GiB'}, no_split_module_classes='CogVLMDecoderLayer')
113
  model = load_checkpoint_and_dispatch(
114
  model,
115
  'local/path/to/hf/version/chat/model', # typical, '~/.cache/huggingface/hub/models--THUDM--cogvlm-chat-hf/snapshots/balabala'
 
109
  low_cpu_mem_usage=True,
110
  trust_remote_code=True,
111
  )
112
+ device_map = infer_auto_device_map(model, max_memory={0:'20GiB',1:'20GiB','cpu':'16GiB'}, no_split_module_classes=['CogVLMDecoderLayer', 'TransformerLayer'])
113
  model = load_checkpoint_and_dispatch(
114
  model,
115
  'local/path/to/hf/version/chat/model', # typical, '~/.cache/huggingface/hub/models--THUDM--cogvlm-chat-hf/snapshots/balabala'