I think in the current scenario I don’t have enough memory if I keep all gradients. I think, however that setting gradients to None well decrease the training speed, because we will need to allocate new memory for gradients after each time we set them to None, Am I correct? In that case, how much do you think this will cost in term of time.
Also her the training loop is very standard, and looks like
        for x, y, indices in iterator:
            x = x.to(self.device).type(torch.float32)
            y = y.to(self.device)

            if self.is_binary_classification:
                y = y.type(torch.float32)

            self.optimizer.zero_grad()

            y_pred = self.model(x).squeeze()

            loss_vec = self.criterion(y_pred, y)
            if weights is not None:
                weights = weights.to(self.device)
                loss = (loss_vec.T @ weights[indices]) / loss_vec.size(0)
            else:
                loss = loss_vec.mean()
            loss.backward()

Rq: criterion is initialized with reduction=None.