@ptrblck . No I am not using DataLoader but maybe I have to change my generator and add length on it!? then using DataLoader? here is the code and I am adding conditioning to VAE
K = int(input("Enter K:"))
# construct the argument parser and parser the arguments
parser = argparse.ArgumentParser(description='VAE Example')
parser.add_argument('--batch-size', type=int, default=2048, metavar='N',
                    help='input batch size for training (default: 128)')
parser.add_argument('--epochs', type=int, default=10, metavar='N',
                    help='number of epochs to train (default: 10)')
parser.add_argument('--seed', type=int, default=1, metavar='S',
                    help='random seed (default: 1)')
parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                    help='how many batches to wait before logging training status')
args = parser.parse_args()

torch.manual_seed(args.seed)

# we use CPU for computation
#device = torch.device("cpu")
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")

def gen_batch(BATCH_SIZE):
    labels = torch.randint(0, 8, (BATCH_SIZE,)).long().to(device)
    theta = (np.pi/4) * labels.float().to(device)
    centers = torch.stack((torch.cos(theta), torch.sin(theta)), dim = -1)
    noise = torch.randn_like(centers) * 0.1

    return centers + noise, labels

def data_gen(BATCH_SIZE):
    #8 gaussians
    while 1:
        yield gen_batch(BATCH_SIZE)

train_loader,train_labels = next(data_gen(args.batch_size))
test_loader = train_loader
cond_dim = train_labels.unique().size(0)

class CVAE(nn.Module):
    def __init__(self, c_dim):
        super(CVAE, self).__init__()
        # encoder part
        self.fc0 = nn.Linear(2 + c_dim, K)
        self.fc1 = nn.Linear(K, K)
        self.fc21 = nn.Linear(K, K)
        self.fc22 = nn.Linear(K, K)
        # decoder part
        self.fc3 = nn.Linear(K + c_dim, K)
        self.fc4 = nn.Linear(K, K)
        self.fc5 = nn.Linear(K, 2)

    def encode(self, x, c):
        concat_input = torch.cat([x, c], 1)
        h1 = F.selu(self.fc1(F.selu(self.fc0(concat_input))))
        return self.fc21(h1), self.fc22(h1)

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5*logvar)
        eps = torch.randn_like(std)
        return mu + eps*std
    
    def sample(self, n=2048):
        """
        Generate n samples from the generative model.
        """
        sample = torch.randn(n, K).to(device)
        out = cvae.decode(sample)

        return out.cpu().detach().numpy()


    def decode(self, z, c):
        concat_input = torch.cat([z, c], 1)
        return self.fc5(F.selu(self.fc4(F.selu(self.fc3(concat_input)))))

    def forward(self, x, c):
        mu, logvar = self.encode(x.view(-1, 2), c)
        z = self.reparameterize(mu, logvar)
        return self.decode(z, c), mu, logvar
cvae = CVAE(c_dim=cond_dim).to(device)
optimizer = optim.Adam(cvae.parameters(), lr=1e-3)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, verbose = True, threshold = 1E-2, eps=1e-6)
# Reconstruction + KL divergence losses summed over all elements and batch
def loss_function(recon_x, x, mu, logvar):
    #BCE = F.binary_cross_entropy(recon_x, x.view(-1, 2), reduction='sum')
    L2 = torch.mean((recon_x-x)**2)
 
    # see Appendix B from VAE paper:
    # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
    # https://arxiv.org/abs/1312.6114
    # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    KLD = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
    return L2 + KLD

def one_hot(labels, class_size): 
    targets = torch.zeros(labels.size(0), class_size)
    for i, label in enumerate(labels):
        targets[i, label] = 1
    return Variable(targets)

def train(epoch):
    cvae.train()
    train_loss = 0
   
    for batch_idx, (data,cond) in enumerate(train_loader):
        if batch_idx > 100:
            break #100 batches per epoch
       # data = data.to(device)
        data, cond = data.to(device), one_hot(cond, cond_dim).to(device)
        optimizer.zero_grad()
        recon_batch, mu, logvar = cvae(data, cond)
        print("mu", mu)
        print("logvar", logvar)
        loss = loss_function(recon_batch, data, mu, logvar)
        recon_batch = recon_batch.detach().numpy()
       
        loss.backward()
        train_loss += loss.item()
        optimizer.step()
    train_loss /= 100
    scheduler.step(train_loss)


    samples_model = cvae.sample()
    samples_gt = gen_batch(2048)     
    kde_skl = KernelDensity(kernel='gaussian', bandwidth=0.2)
    kde_skl.fit(samples_gt)

    # score_samples() returns the log-likelihood of the samples
    log_pdf = kde_skl.score_samples(samples_model)
     
   
    print ("Epoch  average reconstruction error   :" ,epoch ,  train_loss)
    print (" likelihood  :", np.mean(log_pdf))
    return train_loss, np.mean(log_pdf)

if __name__ == "__main__":
    print("Starting training....")
    train_5 = []
    log_5 = []
  
    for epoch in tqdm(range(1, args.epochs + 1)):
     
        a,b= train(epoch)
    #    with open("output.txt","a") as f:
    #        print("epoch:",epoch,",","re:",2+a,",","li:",b,",","los", file=f)
    #        print("VAL ---","re:",2+a,",","li:",b,",","los", file=f)
    #        print("update best loss", file=f)
        
        
  
            
    gt = next(train_loader)
        
    with torch.no_grad():
        out1 = cvae.sample()
        recon = cvae(gt)[0].cpu().numpy()

    rx,ry = recon[:,0], recon[:,1]
        
    gt = gt.cpu().numpy()
    gx, gy = gt[:,0], gt[:,1]
        
    xs, ys = out1[:,0], out1[:,1]

    plt.scatter(gx, gy, c = 'red', s=3)
    plt.scatter(xs, ys, c = 'blue', s=3)
    plt.axes().set_aspect('equal')
    plt.show()