revert from_or_to_torch_tensor.py

Oneflow-Inc · Jun 27, 2024 · b5d2773 · b5d2773
1 parent 62bcb09
commit b5d2773
Showing 1 changed file with 2 additions and 18 deletions.
diff --git a/python/oneflow/utils/tensor/from_or_to_torch_tensor.py b/python/oneflow/utils/tensor/from_or_to_torch_tensor.py
@@ -62,15 +62,7 @@ def from_torch(torch_tensor):
  except:
  print_error_msg()
  assert isinstance(torch_tensor, torch.Tensor)
- # return flow.from_dlpack(torch.to_dlpack(torch_tensor))
- dtype = flow.float16
- if torch_tensor.dtype == torch.int64:
- dtype = flow.int64
- elif torch_tensor.dtype != torch.float16:
- print(torch_tensor.dtype)
- return flow.tensor(
- torch_tensor.cpu().numpy(), device=flow.device("npu"), dtype=dtype
- ).reshape([x for x in torch_tensor.shape])
+ return flow.from_dlpack(torch.to_dlpack(torch_tensor))
 
 
 def to_torch(flow_tensor):
@@ -112,12 +104,4 @@ def to_torch(flow_tensor):
  "WARNING: `to_torch` received a global tensor. A PyTorch CPU tensor which is a copy of its data will be returned."
  )
  return torch.from_numpy(flow_tensor.numpy())
- # return torch.from_dlpack(flow.to_dlpack(flow_tensor))
- dtype = torch.float16
- if flow_tensor.dtype == flow.int64:
- dtype = torch.int64
- elif flow_tensor.dtype != flow.float16:
- print(flow_tensor.dtype)
- return torch.tensor(flow_tensor.numpy(), device="npu", dtype=dtype).reshape(
- [x for x in flow_tensor.shape]
- )
+ return torch.from_dlpack(flow.to_dlpack(flow_tensor))