Why does my training explode?

Question

I'm trying to implement A new lower and upper bound estimation model using gradient descend training method for wind speed interval prediction

For simplicity purposes, I've changed the training data. The network should predict upper and lower bound on when an exponential function would cross a threshold, given a certain number of noisy samples in known intervals.

Here's my implementation:

First, here's the toy data generation method:

def gen_data(n: int) -> (torch.Tensor, torch.Tensor):
    t = torch.linspace(0, 20, NUM_FEATURES).repeat(n, 1)
    a = torch.rand((n, 1))
    lim = torch.max(
        torch.hstack([
            torch.normal(0.2, 0.1, (n, 1)),
            0.01 * torch.ones(n, 1)
        ]),
        -1
    ).values

    x = torch.exp(-a * t) + torch.normal(0, 0.01, t.size())
    y = -torch.log(lim) / a.T
    return x, y[0, :]

Now here's the loss implementation as per the linked article:

def huber(x):
    return torch.where(
        torch.abs(x) < D,
        (x ** 2) / 2,
        D * torch.abs(x) - (D ** 2) / 2
    )


def lube_loss(preds, truth):
    lower, upper = preds.T
    l1 = torch.abs(truth - (lower + upper) / 2)
    l2 = torch.abs(upper - lower)
    l3 = torch.where(
        torch.logical_and(lower < truth, truth < upper),
        torch.zeros_like(truth),
        l1 + l2/2
    )
    return torch.mean(huber(K1 * l1) + huber(K2 * l2 + K3 * l3))

And here's the train code:

def init_weights(m: torch.nn.Module):
    if not hasattr(m, 'weight'):
        return
    n = len(m.weight)
    torch.nn.init.normal_(m.weight, 0, 1 / math.sqrt(n))


def nn_unit(idx: int, inputs: int, outputs: int) -> torch.nn.Module:
    return torch.nn.Sequential(OrderedDict([
        (f'h{idx}', torch.nn.Linear(inputs, outputs)),
        (f'a{idx}', torch.nn.LeakyReLU()),
        (f'drop{idx}', torch.nn.Dropout())]))


def train():
    dataset = torch.utils.data.TensorDataset(*gen_data(100000))
    loader = torch.utils.data.DataLoader(
        dataset, batch_size=BATCH_SIZE, shuffle=True)
    net = torch.nn.Sequential(OrderedDict([
        ('u1', nn_unit(1, NUM_FEATURES, NUM_FEATURES * 4)),
        ('u2', nn_unit(2, NUM_FEATURES * 4, NUM_FEATURES * 8)),
        ('u3', nn_unit(3, NUM_FEATURES * 8, NUM_FEATURES * 4)),
        ('u4', nn_unit(4, NUM_FEATURES * 4, NUM_FEATURES * 4)),
        ('u5', nn_unit(5, NUM_FEATURES * 4, NUM_FEATURES * 2)),
        ('u6', nn_unit(6, NUM_FEATURES * 2, NUM_FEATURES * 2)),
        ('out', torch.nn.Linear(NUM_FEATURES * 2, 2)),
    ]))
    net.apply(init_weights)
    optimizer = torch.optim.SGD(
        net.parameters(), lr=1e-9, momentum=0.5, weight_decay=10
    )
    optimizer.zero_grad()
    running_loss = 0
    for i in range(1000000):
        for x, y in loader:
            preds = net(x)
            loss = lube_loss(preds, y)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        if i % 10 == 9:
            print('[%5d] loss: %.3f' %
                  (i + 1, running_loss / BATCH_SIZE))
            running_loss = 0

    lower, upper = net(X[:10, :]).T
    for l, u, y in zip(lower, upper, Y):
        print(f'{y}: [{l} - {u}]')

    return net

And for completeness the imports and constants:

import math
import torch
import torch.utils.data
from collections import OrderedDict

NUM_FEATURES = 10
K1 = 2
K2 = 1
K3 = 4
D = 0.5
BATCH_SIZE = 100

The train function initially lowers the loss gradually, but then it starts to explode... I've tried lowering the learning rate (started at 0.1 and got to 1e-9), I've tried adding weight decay (started without it, and got up to 10). I've started with a shallow net with just one hidden layer and grew it (hoping it would be able to fit better with more layers and larger layers), it always ended up exploding... What am I missing?