Ok, I’ve got the gradients - I hope - via setting requires_grad to True, and adding hooks at each layer:
class Test_Network(nn.Module):
    def __init__(self):
        super(Test_Network, self).__init__()
        
        self.gradients = None

        #

        layer_list = []

        #

        layer = nn.Linear(4, 6)
        layer_list.append(layer)

        #

        layer = nn.Linear(6, 1)
        layer_list.append(layer)

        #

        self.layer_list = nn.ModuleList(layer_list)


        
    def to(self, device):
        model = super(Test_Network, self).to(device)

        for i, layer in enumerate(model.layer_list):
            model.layer_list[i] = layer.to(device)

        return model


    def save_gradient(self, grad):
        print(grad)

        self.gradients = grad


    def forward(self, x):
        x = x.float()

        for i, layer in enumerate(self.layer_list):
            x.register_hook(self.save_gradient)

            x = layer(x)    

        return x


    

#

def scratch():
    model = Test_Network()

    #

    datum = [0.9,0.1,0.2,0.4]

    batch = np.array([datum])

    batch_tensor = torch.from_numpy(batch)

    batch_tensor.requires_grad = True

    prediction = model(batch_tensor)
    prediction.backward()

    pprint.pprint(batch_tensor.grad)

I’m now a bit lost in trying to transfer these gradients into a heatmap of which nodes of each layer were important to the prediction.