rvae.decode() fails to decode

[utkarshp1161]

Below code expect a 2 dimensional z_mean (like the vanilla vae case)

z_mean, z_std = rvae.encode(norm_patches)
rvae.decode(z_mean)

• when tried with last two latent variables: rvae.decode(z_mean[:, 1:]) -> the decoded patches look very different

Tried a hacky way to decode following: forward_compute_elbo

z_mean, z_std = rvae.encode(norm_patches)
z1, z2, z3 = z_mean[:,0], z_mean[:, 1:3], z_mean[:, 3:]
x_coord_ = rvae.x_coord.expand(norm_patches.shape[0], *rvae.x_coord.size()).cpu()
phi = torch.tensor(z1) # angle
dx = torch.tensor(z2)  # translation
dx = (dx * rvae.dx_prior).unsqueeze(1)
x_coord_ = utils.transform_coordinates(x_coord_, phi, dx)
decoded_patches = rvae.decoder_net(x_coord_.to("cuda"), torch.tensor(z3).to("cuda"))

It will be useful to have this functionality directly using vae.decode(z_mean)

[ziatdinovmax]

Hmm... can you please elaborate on the way in which it failed to decode? Did it throw an error?

[utkarshp1161]

Sure. Here is the link to elaborated error msg.

[ziatdinovmax]

@utkarshp1161 - Ah OK, I see what you mean. We can re-write the decode method for rvae, based on the snippet you provided earlier. Would you like to give it a try?

[utkarshp1161]

@ziatdinovmax, sure.

[utkarshp1161]

Raised a pull request here #90