i set the Batch size=1;but the CUDA out of memory…I don’t know what… this is my train

username: yi_zhuang
cooked: <pre><code>def train(self):

    img, label, test_data, test_label = self.get_source_data()

    img = torch.from_numpy(img)
    label = torch.from_numpy(label - 1)

    dataset = torch.utils.data.TensorDataset(img, label)
    self.dataloader = torch.utils.data.DataLoader(dataset=dataset, batch_size=self.batch_size, shuffle=True)

    test_data = torch.from_numpy(test_data)
    test_label = torch.from_numpy(test_label - 1)
    test_dataset = torch.utils.data.TensorDataset(test_data, test_label)
    self.test_dataloader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=self.batch_size, shuffle=True)

    # Optimizers
    self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr, betas=(self.b1, self.b2))

    test_data = Variable(test_data.type(self.Tensor))
    test_label = Variable(test_label.type(self.LongTensor))

    bestAcc = 0
    averAcc = 0
    num = 0
    Y_true = 0
    Y_pred = 0

    # Train the cnn model
    total_step = len(self.dataloader)
    curr_lr = self.lr

    for e in range(self.n_epochs):
        # in_epoch = time.time()
        self.model.train()
        for i, (img, label) in enumerate(self.dataloader):

            img = Variable(img.cuda().type(self.Tensor))
            label = Variable(label.cuda().type(self.LongTensor))

            # data augmentation
            aug_data, aug_label = self.interaug(self.allData, self.allLabel)
            img = torch.cat((img, aug_data))
            label = torch.cat((label, aug_label))


            tok, outputs = self.model(img)

            loss = self.criterion_cls(outputs, label) 

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()


        # out_epoch = time.time()


        # test process
        if (e + 1) % 1 == 0:
            self.model.eval()
            Tok, Cls = self.model(test_data)


            loss_test = self.criterion_cls(Cls, test_label)
            y_pred = torch.max(Cls, 1)[1]
            acc = float((y_pred == test_label).cpu().numpy().astype(int).sum()) / float(test_label.size(0))
            train_pred = torch.max(outputs, 1)[1]
            train_acc = float((train_pred == label).cpu().numpy().astype(int).sum()) / float(label.size(0))

            print('Epoch:', e,
                  '  Train loss: %.6f' % loss.detach().cpu().numpy(),
                  '  Test loss: %.6f' % loss_test.detach().cpu().numpy(),
                  '  Train accuracy %.6f' % train_acc,
                  '  Test accuracy is %.6f' % acc)

            self.log_write.write(str(e) + "    " + str(acc) + "\n")
            num = num + 1
            averAcc = averAcc + acc
            if acc &gt; bestAcc:
                bestAcc = acc
                Y_true = test_label
                Y_pred = y_pred


    torch.save(self.model.module.state_dict(), 'model.pth')
    averAcc = averAcc / num
    print('The average accuracy is:', averAcc)
    print('The best accuracy is:', bestAcc)
    self.log_write.write('The average accuracy is: ' + str(averAcc) + "\n")
    self.log_write.write('The best accuracy is: ' + str(bestAcc) + "\n")

    return bestAcc, averAcc, Y_true, Y_pred
    # writer.close()
</code></pre>