Here is my training loop for the classification:
print("device = ", "cuda" if torch.cuda.is_available() else "cpu")
in_size = (128, 248)
trainDataset = SimpleDataSet(root_data_dir=root_data_path, samples_size=in_size)
trainLoader = torch.utils.data.DataLoader(trainDataset, batch_size=50, shuffle=True,
                                            num_workers=0, collate_fn=trainDataset.collate_fn)
margin = 1.0
net_classification = ClassificationNet().train().to(device)
criterion = nn.CrossEntropyLoss(reduction='mean')
lr1 = 0.0001
optimizer = optim.Adam(net_classification.parameters(), lr=lr1, weight_decay=0.0005)
starting_epoch = 0
for epoch in range(starting_epoch,1000):  # loop over the dataset multiple times
    print("Epoch %d started!"%(epoch+1))
    running_loss = 0.0
    epoch_loss = 0.0
    running_acc = 0.0
    epoch_acc = 0.0
    count = 0.0
    epoch_count = 0.0
    for i, (samples, ids) in enumerate(trainLoader):
        samples = samples.to(device)
        ids = ids.to(device)

        optimizer.zero_grad()
        
        out = net_classification(samples)
        loss = criterion(out ,ids)
        loss.backward()
        optimizer.step()
        acc = net_classification.calcAcc(preds=out, targets=ids)
        
        running_loss += loss.item()
        epoch_loss += loss.item()*samples.size(0)
        running_acc += acc
        epoch_acc += acc*samples.size(0)
        count += 1
        epoch_count += samples.size(0)
        if i % 2 == 1:
            print('[%d, %5d] loss: %.5f' %
                (epoch + 1, i + 1, running_loss / count))
            print("[%d, %5d] acc: %.5f" % 
                (epoch + 1, i+1, running_acc / count))
            running_loss = 0.0
            running_acc = 0.0
            count = 0.0
    print(("Epoch %d finished!") % (epoch+1))
    print("Epoch loss: ", epoch_loss / epoch_count)
    print("Epoch acc: ", epoch_acc / epoch_count)
    if epoch % 10 == 0:
        torch.save({
            'epoch': epoch,
            'model_state_dict': net_classification.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': epoch_loss / epoch_count,
            }, out_weights_dir+"voice_net5_"+str(epoch+1)+".pth")

print('Finished Training')

And this is my training loop for embedding generation:
print("device = ", "cuda" if torch.cuda.is_available() else "cpu")
trainDataset = BalancedDataset(root_data_path=root_data_path, k=4)
trainLoader = torch.utils.data.DataLoader(trainDataset, batch_size=5, shuffle=True,
                                            num_workers=0, collate_fn=trainDataset.collate_fn)
classification_net = ClassificationNet()
checkpoint = torch.load("./weights/class_net_101.pth")
classification_net.load_state_dict(checkpoint['model_state_dict'])

margin = 1.0
lr1 = 0.0001  
net = Net().train()
net.features.load_state_dict(classification_net.features.state_dict())
net = net.to(device)
del checkpoint
del classification_net
optimizer = optim.Adam(net.parameters(), lr=lr1, weight_decay=0.0005)
criterion = OnlineTripletLoss(margin, HardestNegativeTripletSelector(margin))
starting_epoch = 0
for epoch in range(starting_epoch,1000):  # loop over the dataset multiple times
    print("Epoch %d started!"%(epoch+1))
    running_loss = 0.0
    epoch_loss = 0.0
    count = 0.0
    epoch_count = 0.0
    for i, (x, ids) in enumerate(trainLoader):
        x = x.to(device)
        ids = ids.to(device)
        # zero the parameter gradients
        optimizer.zero_grad()
        # forward + backward + optimize
        features = net(x)
        
        loss, n_triplets = criterion(features, ids)
        assert n_triplets > 0
        loss.backward()
        optimizer.step()
        # print statistics
        running_loss += loss.item()*n_triplets
        epoch_loss += loss.item()*n_triplets
        count += n_triplets
        epoch_count += n_triplets
        if i % 2 == 1:
            print('[%d, %5d] loss: %.5f' %
                (epoch + 1, i + 1, running_loss / count))
            running_loss = 0.0
            count = 0.0
    print(("Epoch %d finished!") % (epoch+1))
    print("Epoch loss: ", epoch_loss / epoch_count)
    if epoch % 10 == 0:
        torch.save({
            'epoch': epoch,
            'model_state_dict': net.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': epoch_loss / epoch_count,
            }, out_weights_dir+"net4_"+str(epoch+1)+".pth")
print('Finished Training')