remove global default tensor when using cuda
@@ -113,13 +113,13 @@ class SSVAE(nn.Module):
with pyro.iarange("data"):
# sample the handwriting style from the constant prior distribution
- prior_loc = torch.zeros([batch_size, self.z_dim])
- prior_scale = torch.ones([batch_size, self.z_dim])
+ prior_loc = xs.new_zeros([batch_size, self.z_dim])
+ prior_scale = xs.new_ones([batch_size, self.z_dim])
zs = pyro.sample("z", dist.Normal(prior_loc, prior_scale).independent(1))
# if the label y (which digit to write) is supervised, sample from the
# constant prior, otherwise, observe the value (i.e. score it against the constant prior)
- alpha_prior = torch.ones([batch_size, self.output_size]) / (1.0 * self.output_size)
+ alpha_prior = xs.new_ones([batch_size, self.output_size]) / (1.0 * self.output_size)
ys = pyro.sample("y", dist.OneHotCategorical(alpha_prior), obs=ys)
# finally, score the image (x) using the handwriting style (z) and
@@ -173,7 +173,7 @@ class SSVAE(nn.Module):
res, ind = torch.topk(alpha, 1)
# convert the digit(s) to one-hot tensor(s)
- ys = torch.zeros(alpha.size())
+ ys = xs.new_zeros(alpha.size())
ys = ys.scatter_(1, ind, 1.0)
return ys
@@ -205,8 +205,8 @@ class SSVAE(nn.Module):
"""
# sample the handwriting style from the constant prior distribution
- prior_loc = torch.zeros([batch_size, self.z_dim])
- prior_scale = torch.ones([batch_size, self.z_dim])
+ prior_loc = ys.new_zeros([batch_size, self.z_dim])
+ prior_scale = ys.new_ones([batch_size, self.z_dim])
zs = pyro.sample("z", dist.Normal(prior_loc, prior_scale).independent(1))
# sample an image using the decoder
@@ -299,9 +299,6 @@ def main(args):
:param args: arguments for SS-VAE
:return: None
"""
- if args.cuda:
- torch.set_default_tensor_type('torch.cuda.FloatTensor')
-
if args.seed is not None:
set_seed(args.seed, args.cuda)
