Skip to content

GitLab

Unverified Commit c198c7c0 authored by Super Daniel's avatar Super Daniel Committed by GitHub
Browse files

[hotfix] meta tensor default device. (#2510)

parent 077a5cdd
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
colossalai/fx/profiler/tensor.py
View file @ c198c7c0
...@@ -43,7 +43,7 @@ class MetaTensor(torch.Tensor): ...@@ -43,7 +43,7 @@ class MetaTensor(torch.Tensor):
storage_offset=elem.storage_offset(), storage_offset=elem.storage_offset(),
dtype=elem.dtype, dtype=elem.dtype,
layout=elem.layout, layout=elem.layout,
device=fake_device if fake_device is not None else torch.device('cpu'), device=fake_device or (elem.device if elem.device.type != 'meta' else torch.device('cpu')),
requires_grad=elem.requires_grad) # deceive the frontend for aten selections requires_grad=elem.requires_grad) # deceive the frontend for aten selections
r._tensor = elem r._tensor = elem
# ...the real tensor is held as an element on the tensor. # ...the real tensor is held as an element on the tensor.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment