Thanks for the code!
I’ve let the model train for a few hours using PyTorch 1.9.1+cu113 as well as 1.8.2+cu111 on an A6000 using this code:
import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader
import numpy as np


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        #self.conv1 = nn.Conv2d(3, 6, 5)
        #self.fc1 = nn.Linear(72*12, 768)
        #self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(2*3*64*64, 80*80)
        self.fc2 = nn.Linear(80*80, 125*125)
        #self.fc3 = nn.Linear(2*3*16*16, 60*60)
        self.us = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.conv1 = nn.Conv2d(1, 64, 5,padding=2) # would decrease size 4 in each dimension
        self.conv2 = nn.Conv2d(64, 32, 5,padding=2)
        self.conv3 = nn.Conv2d(32, 16, 5,padding=2)
        self.conv4 = nn.Conv2d(16, 4, 5,padding=2)
        self.conv5 = nn.Conv2d(4, 1, 5,padding=2)


    # With relu and wihout relu has similar result    
    def forward(self, x):
        x = x.view(-1, 2*3*64*64)
        #x = self.conv1(x)
        #x = x.view(-1, 72*12)
        ## Fully connected layers
        x = self.fc1(x)
        #x = F.relu(self.fc1(x))
        x = nn.LeakyReLU(0.1)(x)
        x = self.fc2(x)
        # x = F.relu(self.fc2(x))
        x = nn.LeakyReLU(0.1)(x)	
        x = torch.reshape(x, (np.shape(x)[0], -1, 125, 125)) # 
        ## Upsampling
        x = self.us(x)
        ## Five convolution layers
        x = nn.LeakyReLU(0.1)( self.conv1(x) )
        x = nn.LeakyReLU(0.1)( self.conv2(x) )
        x = nn.LeakyReLU(0.1)( self.conv3(x) )
        x = nn.LeakyReLU(0.1)( self.conv4(x) )
        x = self.conv5(x)
        x = x.view(np.shape(x)[0], -1)
        # x = self.conv1(x)
        return x
#Generate model
net = Net().double()
#Batch size too small, no parallelization, only 1 GPU
#device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.is_available():
    device = torch.device("cuda:0")
    print("Use GPU")
else:
    device = torch.device("cpu")
    print("Use CPU")


x = torch.randn(64*100, 2, 3, 64, 64).double()
y = torch.randn(64*100, 62500).double() 

dataset = TensorDataset(x, y)
trainloader = DataLoader(dataset, batch_size=64)
criterion = nn.MSELoss()

#Use all 3 GPUs
# if torch.cuda.device_count() > 1:
#     print("Let's use", torch.cuda.device_count(), "GPUs!")
#     net = nn.DataParallel(net)
#Put data on all GPUs

net.to(device)
#Learning Process
# Train the network
nepoch = 4
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
for epoch in range(nepoch):
    print('Epoch: %d\n' %(epoch))
    for i, data in enumerate(trainloader, 0):
        #print(i)
        inputs, labels = data[0].to(device), data[1].to(device)
        # Zero gradient
        optimizer.zero_grad()
        # Forward, backward(gradient), optimize (add gradient to weight)
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        # Print loss information
        print(loss.item())

and could reproduce any issue.