'meta' device instead of cuda
Browse files
modeling_InternLM_XComposer.py
CHANGED
@@ -67,7 +67,7 @@ class InternLMXComposerForCausalLM(PreTrainedModel):
|
|
67 |
# speed up init llm
|
68 |
with torch.device('meta'):
|
69 |
self.internlm_model = InternLMForCausalLM._from_config(config)
|
70 |
-
|
71 |
|
72 |
self.internlm_proj = nn.Linear(4096,
|
73 |
self.internlm_model.config.hidden_size)
|
|
|
67 |
# speed up init llm
|
68 |
with torch.device('meta'):
|
69 |
self.internlm_model = InternLMForCausalLM._from_config(config)
|
70 |
+
self.internlm_model.to_empty(device=config.device).to(torch.float16)
|
71 |
|
72 |
self.internlm_proj = nn.Linear(4096,
|
73 |
self.internlm_model.config.hidden_size)
|