oneflow
oneflow copied to clipboard
Fix div dtype & atan2
@BBuf fix https://github.com/Oneflow-Inc/oneflow/issues/9154 & https://github.com/Oneflow-Inc/OneTeam/issues/1652
div的问题其实就是设置scalar_div的输出是float就行,atan2的问题就在于在输入为只有一个元素的张量时需要广播一下(多个维度的时候torch不支持广播)
import torch as flow
# import oneflow as flow
import numpy as np
x = flow.tensor([[1], [1]]).to("cuda")
# x = flow.ones((2,3,4)).to("cuda")
arr = np.arange(24).reshape(2,3,4).astype(np.float32)
y = flow.tensor(arr).to("cuda")
print(x.shape, y.shape)
z = flow.atan2(x, y)
z1 = flow.atan2(y, x)
print("z >>>>>>>>>>>>>> type:%s;\n%s" % (type(z), z))
print("z >>>>>>>>>>>>>> type:%s;\n%s" % (type(z1), z1))
RuntimeError: The size of tensor a (2) must match the size of tensor b (3) at non-singleton dimension 1