args.rec_hidden,
128,
learn_emb=args.learn_emb).to(device)
elif args.enc == 'mtan_rnn':
rec = models.enc_mtan_rnn(dim,
torch.linspace(0, 1., args.num_ref_points),
args.latent_dim,
args.rec_hidden,
embed_time=128,
learn_emb=args.learn_emb,
num_heads=args.enc_num_heads).to(device)
if args.dec == 'rnn3':
dec = models.dec_rnn3(dim,
torch.linspace(0, 1., 128),
args.latent_dim,
args.gen_hidden,
128,
learn_emb=args.learn_emb).to(device)
elif args.dec == 'mtan_rnn':
dec = models.dec_mtan_rnn(dim,
torch.linspace(0, 1., args.num_ref_points),
args.latent_dim,
args.gen_hidden,
embed_time=128,
learn_emb=args.learn_emb,
num_heads=args.dec_num_heads).to(device)
classifier = models.create_classifier(args.latent_dim,
args.rec_hidden).to(device)
params = (list(rec.parameters()) + list(dec.parameters()) +
list(classifier.parameters()))
elif args.enc == "mtan_rnn":
rec = models.enc_mtan_rnn(
dim,
torch.linspace(0, 1.0, args.num_ref_points),
args.latent_dim,
args.rec_hidden,
embed_time=128,
learn_emb=args.learn_emb,
num_heads=args.enc_num_heads,
).to(device)
if args.dec == "rnn3":
dec = models.dec_rnn3(
dim,
torch.linspace(0, 1.0, args.num_ref_points),
args.latent_dim,
args.gen_hidden,
128,
learn_emb=args.learn_emb,
).to(device)
elif args.dec == "mtan_rnn":
dec = models.dec_mtan_rnn(
dim,
torch.linspace(0, 1.0, args.num_ref_points),
args.latent_dim,
args.gen_hidden,
embed_time=128,
learn_emb=args.learn_emb,
num_heads=args.dec_num_heads,
).to(device)
params = list(dec.parameters()) + list(rec.parameters())