BIAS_TRAIN = (train_data_size - 1) / (args.batch_size - 1) BIAS_TEST = (test_data.dataset.__len__() - 1) / (args.batch_size - 1) def cuda_tensors(obj): for attr in dir(obj): value = getattr(obj, attr) if isinstance(value, torch.Tensor): setattr(obj, attr, value.cuda()) encA = EncoderA(args.wseed, zPrivate_dim=args.n_private) decA = DecoderA(args.wseed, zPrivate_dim=args.n_private) encB = EncoderB(args.wseed) decB = DecoderB(args.wseed) if CUDA: encA.cuda() decA.cuda() encB.cuda() decB.cuda() cuda_tensors(encA) cuda_tensors(decA) cuda_tensors(encB) cuda_tensors(decB) optimizer = torch.optim.Adam(list(encB.parameters()) + list(decB.parameters()) + list(encA.parameters()) + list(decA.parameters()), lr=args.lr)