Views
No views yet
PatchEmbed + Transformer blocks + fixed sin-cos positional embeddingmu/logvar + reparameterization1import torch
2from tactile_vae.model import TactileVAE, VAELoss
3
4model = TactileVAE()
5out = model(torch.randn(2, 3, 128, 128))
6loss_fn = VAELoss(beta=1.0)
7losses = loss_fn(out["x_hat"], torch.randn(2, 3, 128, 128), out["mu"], out["logvar"])