Spaces:
Running
on
Zero
Running
on
Zero
Update model.py
Browse files
model.py
CHANGED
@@ -355,6 +355,8 @@ class VoilaModel(LlamaPreTrainedModel):
|
|
355 |
if self.training or \
|
356 |
(past_key_values is None and ref_embs is not None) or \
|
357 |
(past_key_values is not None and past_key_values.get_seq_length() < 4 and ref_embs is not None):
|
|
|
|
|
358 |
ref_embs = self.ref_emb_linear(ref_embs.to(self.ref_emb_linear.weight.dtype))
|
359 |
ref_embs = ref_embs * ref_embs_mask.unsqueeze(-1).unsqueeze(-1)
|
360 |
# (padding_left,padding_right,padding_top,padding_bottom,padding_front,padding_back)
|
|
|
355 |
if self.training or \
|
356 |
(past_key_values is None and ref_embs is not None) or \
|
357 |
(past_key_values is not None and past_key_values.get_seq_length() < 4 and ref_embs is not None):
|
358 |
+
print("model device",self.ref_emb_linear.weight.device)
|
359 |
+
print('tensor device', ref_embs.device)
|
360 |
ref_embs = self.ref_emb_linear(ref_embs.to(self.ref_emb_linear.weight.dtype))
|
361 |
ref_embs = ref_embs * ref_embs_mask.unsqueeze(-1).unsqueeze(-1)
|
362 |
# (padding_left,padding_right,padding_top,padding_bottom,padding_front,padding_back)
|