def _graph_constant(g, value, dims, type, *args, **kwargs):
assert isinstance(value, numbers.Number)
assert type is not None
isscalar = False
if dims is None or dims == 0 or set(dims) == set([0]):
dims = [1]
isscalar = True
type = type.lower()
if type == "char":
tensor = torch.CharTensor(*dims)
elif type == "short":
tensor = torch.ShortTensor(*dims)
elif type == "int":
tensor = torch.IntTensor(*dims)
elif type == "long":
tensor = torch.LongTensor(*dims)
elif type == "half":
tensor = torch.HalfTensor(*dims)
elif type == "float":
tensor = torch.FloatTensor(*dims)
elif type == "double":
tensor = torch.DoubleTensor(*dims)
else:
raise ValueError("Unknown type, type should be one of the following strings: "
"char, short, int, long, half, float, double")
tensor.fill_(value)
if isscalar:
return g.op("Constant", *args, value_z=tensor, **kwargs)
return g.op("Constant", *args, value_t=tensor, **kwargs)
评论列表
文章目录