one_hot variable/tensor bug
@@ -100,6 +100,10 @@ def log_gamma(xx):
def to_one_hot(x, ps):
+ if isinstance(x, Variable):
+ ttype = x.data.type()
+ elif isinstance(x, torch.Tensor):
+ ttype = x.type()
batch_size = x.size(0)
classes = ps.size(1)
# create an empty array for one-hots
@@ -107,7 +111,7 @@ def to_one_hot(x, ps):
# this operation writes ones where needed
batch_one_hot.scatter_(1, x.data.view(-1, 1), 1)
- return Variable(batch_one_hot.type_as(x))
+ return Variable(batch_one_hot.type(ttype))
def tensor_histogram(ps, vs):
