I am working on making distributed training using Dask 3. So far, I have been able to use dask concurrent.futures interface to do training with success. What I basically do is:

Create a list with batches of data points that I need for training the model.
Call a  train_batches  function using concurrent.futures with a distributed scheduler.
Accumulate the gradients returned by the futures, wait for them, and gather them.

Set the gradients 2.
And finally, call  optimizer.step() .

I think I can go further with the distribution of calculations on my code because forward() is receiving a dictionary and I have to iterate over it:
    def forward(self, X):
        """Forward propagation
        This is forward propagation and it returns the atomic energy.
        Parameters
        ----------
        X : list
            List of inputs in the feature space.
        Returns
        -------
        outputs : tensor
            A list of tensors with energies per image.
        """

        outputs = []

        for hash in X:
            image = X[hash]
            atomic_energies = []

            for symbol, x in image:
                x = self.linears[symbol](x)

                intercept_name = 'intercept_' + symbol
                slope_name = 'slope_' + symbol
                slope = getattr(self, slope_name)
                intercept = getattr(self, intercept_name)

                x = (slope * x) + intercept
                atomic_energies.append(x)

            atomic_energies = torch.cat(atomic_energies)
            image_energy = torch.sum(atomic_energies)
            outputs.append(image_energy)
        outputs = torch.stack(outputs)
        return outputs

I am able to get futures that compute the outputs without problems. However, when calling loss,  loss.backard() and doing this loop (that works when I don’t mess with forward()):
for param in model.parameters():
    gradients.append(param.grad.detach().numpy())

it raises an error because the model gradients are None. What would be the reason why the model is not getting the gradients? I would appreciate any discussion on this.