n_categories = 5
batch_size = 32
input_size = 2
hidden_size = 2
lstm_layers = 50

model = LSTMClassifier(n_categories, batch_size, input_size, hidden_size, lstm_layers, dropout=0.5).cuda()

weightSampler = torch.utils.data.WeightedRandomSampler(weights=bfp_data_setup.weights, num_samples=len(bfp_data_setup.weights), replacement=True)

train_dataloader = DataLoader(bfp_data_setup.dataset, batch_size=batch_size, sampler=weightSampler)

loss_fn = torch.nn.CrossEntropyLoss(size_average=False)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)

for t in range(num_epochs):    
    for batch_padded_sequences, batch_packed_lengths, batch_packed_categories in train_dataloader:
        
        cuda_batch_padded_seq = batch_padded_sequences.cuda()
        cuda_batch_padded_lens = batch_packed_lengths.cuda()
        cuda_batch_packed_categories = batch_packed_categories.cuda()
        
        hidden = model.init_hidden(batch_packed_lengths.size(0))
        
        optimizer.zero_grad()
        model.zero_grad()
        
        decoded, output, hidden = model(cuda_batch_padded_seq, cuda_batch_padded_lens, hidden)
        
        loss = loss_fn(decoded, cuda_batch_packed_categories)
        
        epochLoss = loss.item()
        
        print("Epoch {}, Adam {}".format(t, epochLoss))

        loss.backward()
        
        optimizer.step()