When I am training a cbow words_embedding on one GPU, a segmentation error occurs just after one training epoch is finished.
The training code is below:
loss_function=nn.NLLLoss()
optimizer=optim.RMSprop(cbow_embedder.parameters(), lr=0.05) #lstm optimizers, may need to change the lr
cbow_dataset=torch.utils.data.TensorDataset(contexts, labels)
cbow_dataloader=torch.utils.data.DataLoader(cbow_dataset, batch_size=1000, shuffle=False, pin_memory=False, num_workers=4)
for epoch in range(100):
        total_loss=0.0
        recent_loss=0.0
        num_batches=0.0
        ctime=time.time()
        for contexts, labels in cbow_dataloader:
            cbow_embedder.zero_grad()
            log_probs=cbow_embedder(contexts)
            loss=loss_function(log_probs.view((contexts.shape[0], -1)), labels.cuda())
            print(loss)
            loss.backward()
            optimizer.step()
            total_loss+=loss
            recent_loss+=loss
            num_batches+=1.0
            if num_batches%10==0:
                ntime=time.time()
                print(num_batches, list(contexts.size())[0], recent_loss/10)
                ctime=time.time()
                recent_loss=0.0
            if num_batches%500==0:
                torch.save(cbow_embedder, save_dir+'/cbow_embedding')
                print('saved cbow')
        print('epoch', epoch, 'avg_loss',  total_loss/num_batches)

Can anyone tell what has happened here? Thanks!