def setup_tvae():
enc = TVAEEncoder([2, 1], nonlinearity=nn.ReLU)
dec = MLP([1, 2], initial_batchnorm=False, nonlinearity=nn.ReLU)
return dt.decomposition.deep.TVAE(enc, dec, learning_rate=1e-3)
def test_mlp_sanity():
mlp = MLP([100, 10, 2])
with torch.no_grad():
x = torch.empty((5, 100)).normal_()
mlp(x)
def setup_tae():
enc = MLP([2, 1], initial_batchnorm=False, nonlinearity=nn.Tanh)
dec = MLP([1, 2], initial_batchnorm=False, nonlinearity=nn.Tanh)
return dt.decomposition.deep.TAE(enc, dec, learning_rate=1e-3)