for epoch in range(n_epochs):
    model.train()
    tloss = []
    for b in range(train_data.shape[0] // bsize):
        tx = torch.from_numpy(train_data[b*bsize:b*bsize+bsize]).to(device)
        tx_hat = model(tx,ts,tD,tU) # the input (tx) is of size (bsize,5000,3), and output is the same
        loss = loss_fn(tx,tx_hat)
        optim.zero_grad()
        loss.backward()
        optim.step()
        scheduler.step()
        tloss.append(loss.item())
        
    # validate
    model.eval()
    vloss = []
    with torch.no_grad(): 
        for b in range(val_data.shape[0]//bsize):
            tx = torch.from_numpy(train_data[b*bsize:b*bsize+bsize]).to(device)
            tx_hat = model(tx,ts,tD,tU)
            loss = loss_fn(tx,tx_hat)
            vloss.append(loss.item())

It started working fine when I wrote my own L1 loss
loss_fn = torch.abs(out - targets).sum() / (out.size(0)*out.size(1)*out.size(2))
I also have never seen this problem before and when I just generated random tensors of the same size and applied the L1Loss() function I get sensible answers. And the fact that it works correctly torch.no_grad() scope makes me believe that it must be some weird bug. My network is kinda of complicated and involves a lot of slicing, however I can’t share the details of it publicly, sorry.
Edit: I also used the same model in python3 (but on a different dataset) and I was getting sensible outputs, so I think it might be python2 specific.