def backward_theta(self, x, y):
        self.vgg16.train()
        loss, prob_trees = self.get_loss(x, y)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

This code above can run normally, but the code below raise the error “CUDA out of memory”. Is something wrong in the else place?
def backward_theta(self, x, y):
        self.vgg16.train()
        loss, prob_trees = self.get_loss(x, y)
        self.optimizer.zero_grad()
        #loss.backward()
        self.optimizer.step()