How can I visualize what happens during loss.backward()?
I am confident in my understanding of the forward pass of my model, how can I control its backward pass?

This is not a theoretical question about what back-propagation is. The question is a practical one, about whether or not there are tools suited to visualize/track/control what happens during back-propagation.

Ideally, this tool would allow to visualize the structure of the computational graph of the model (a graph of the model's operations), its inputs and its trainable parameters.

Now, I do:

loss.backward()


and I would like to visualize what happens in that step.

There has been been already mention of pytorchviz which lets you visualize the graph. 

Here is a small example that might help you to understand how pytorchviz does trace the graph using the grad_fn:

import torch
from torch import nn
d = 5
x = torch.rand(d, requires_grad=True)
print('Tensor x:', x)
y = torch.ones(d, requires_grad=True)
print('Tensor y:', y)
loss = torch.sum(x*y)*3

del x
print()
print('Tracing back tensors:')
def getBack(var_grad_fn):
    print(var_grad_fn)
    for n in var_grad_fn.next_functions:
        if n[0]:
            try:
                tensor = getattr(n[0], 'variable')
                print(n[0])
                print('Tensor with grad found:', tensor)
                print(' - gradient:', tensor.grad)
                print()
            except AttributeError as e:
                getBack(n[0])

loss.backward()
getBack(loss.grad_fn)


Output:

Tensor x: tensor([0.0042, 0.5376, 0.7436, 0.2737, 0.4848], requires_grad=True)
Tensor y: tensor([1., 1., 1., 1., 1.], requires_grad=True)

Tracing back tensors:
&lt;MulBackward object at 0x1201bada0&gt;
&lt;SumBackward0 object at 0x1201bacf8&gt;
&lt;ThMulBackward object at 0x1201bae48&gt;
&lt;AccumulateGrad object at 0x1201badd8&gt;
Tensor with grad found: tensor([0.0042, 0.5376, 0.7436, 0.2737, 0.4848], requires_grad=True)
 - gradient: tensor([3., 3., 3., 3., 3.])

&lt;AccumulateGrad object at 0x1201bad68&gt;
Tensor with grad found: tensor([1., 1., 1., 1., 1.], requires_grad=True)
 - gradient: tensor([0.0125, 1.6129, 2.2307, 0.8211, 1.4543])


Further you should definately take a look into how autograd functions (that are used by the backward()-function) are actually work !
Here is a tutorial from the pytorch site with an easy and short example:

PyTorch: Defining New autograd Functions

Hope this helps a bit!
