I am working on an encoder-decoder architecture to perform regression for a family of sinusoidal functions. For training, my encoder takes in a random subset of input training pairs (total pairs = 40 for each function) and produces a corresponding feature representation (mean averaged over all chosen subset pairs). Then, the decoder takes this feature representation and all the input values to make predict.
This is my initial model. It may have mistakes, but my inital concern is that Nan values are outputted after a certain point from encoder. I have not been able to figure out why. I have checked that my inputs do not have Nan and the learning rate is appropriate.
Nf = 2000 # the number of different functions f that we will generate
Npts = 40 # the number of x values that we will use to generate each fa
x = torch.zeros(Nf, Npts, 1)
for k in range(Nf):
    x[k,:,0] = torch.linspace(-2, 2, Npts)

x += torch.rand_like(x)*0.1
a = -2 + 4*torch.rand(Nf).view(-1,1).repeat(1, Npts).unsqueeze(2)
y = a*torch.sin(x+a)

LR = 0.0001
MAX_EPOCH = 20
BATCH_SIZE = 80

#  use gpu if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

X_train, X_val, y_train, y_val = map(torch.tensor, train_test_split(x, y, test_size=0.2))

train_dataloader = DataLoader(TensorDataset(X_train.squeeze(-1), y_train.squeeze(-1)), batch_size=BATCH_SIZE,
                              pin_memory=True, shuffle=True)
val_dataloader = DataLoader(TensorDataset(X_val.squeeze(-1), y_val.squeeze(-1)), batch_size=BATCH_SIZE,
                            pin_memory=True, shuffle=True)

class Encoder(nn.Module):
  def __init__(self, num_inputs, num_hidden_1, num_hidden_2, num_outputs):
    super(Encoder, self).__init__()
    self.encoder_hidden_layer_1 = nn.Linear(in_features=num_inputs, out_features=num_hidden_1)
    self.encoder_hidden_layer_2 = nn.Linear(in_features=num_hidden_1, out_features=num_hidden_2)
    self.encoder_output_layer = nn.Linear(in_features=num_hidden_2, out_features=num_outputs)

  def forward(self, x):
    x.to(device)
    x1 = F.relu(self.encoder_hidden_layer_1(x))
    x2 = F.relu(self.encoder_hidden_layer_2(x1))
    feature_representation = self.encoder_output_layer(x2)
    
    return feature_representation

class Decoder(nn.Module):
  def __init__(self, num_inputs, num_hidden_1, num_hidden_2, num_outputs):
    super(Decoder, self).__init__()

    self.decoder_hidden_layer_1 = nn.Linear(in_features=num_inputs, out_features=num_hidden_1)
    self.decoder_hidden_layer_2 = nn.Linear(in_features=num_hidden_1, out_features=num_hidden_2)
    self.decoder_output_layer = nn.Linear(in_features=num_hidden_2, out_features=num_outputs)

  def forward(self, x):
    x.to(device)
    x1 = F.relu(self.decoder_hidden_layer_1(x))
    x2 = F.relu(self.decoder_hidden_layer_2(x1))
    y_pred = self.decoder_output_layer(x2)

    return y_pred 

class EncoderDecoder(nn.Module):
  def __init__(self, encoder, decoder, enc_outputs):
    
    super(EncoderDecoder, self).__init__()
    self.num_outputs = enc_outputs
    self.encoder = encoder
    self.decoder = decoder

  def forward(self, x_c, y_c, x_t):
    x_c.to(device)
    y_c.to(device)
    x_t.to(device)

    r_c = torch.empty(1, self.num_outputs).to(device)
    for x, y in zip(torch.transpose(x_c,0,1), torch.transpose(y_c,0,1)):
        enc_input = torch.stack([x,y])
        enc_input = enc_input.view(-1,torch.numel(enc_input))
        enc_output = self.encoder(enc_input)

        r_c = torch.add(r_c, enc_output)

    r_c = torch.div(r_c,torch.numel(x_c))
    dec_input = torch.cat((x_t, r_c),-1)
    dec_input = dec_input.view(-1,torch.numel(dec_input))
    output = self.decoder(dec_input)

    return output

# Model instantiation and initialisation 
def init_weights(m):
  if type(m) == nn.Linear:
      torch.nn.init.xavier_uniform_(m.weight)
      m.bias.data.fill_(0.01)

# Create and initialize encoder:
enc_inputs, enc_hidden_1, enc_hidden_2, enc_outputs = 2, 16, 16, 8
encoder = Encoder(enc_inputs, enc_hidden_1, enc_hidden_2, enc_outputs).to(device)
encoder.apply(init_weights)

# Create and initialize decoder:
dec_inputs, dec_hidden_1, dec_hidden_2, dec_outputs = 9, 32, 16, 1
decoder = Decoder(dec_inputs, dec_hidden_1, dec_hidden_2, dec_outputs).to(device)
decoder.apply(init_weights)

# Create and initialize EncoderDecoder:
model = EncoderDecoder(encoder,decoder,enc_outputs).to(device)

# Adam optimizer with learning rate LR
optimizer = optim.Adam(model.parameters(), lr=LR)

# mean-squared error loss
criterion = nn.MSELoss()

def train(model, train_dataloader, optimizer, criterion, MAX_EPOCH):
  for epoch in range(MAX_EPOCH):
    model.train()
    loss = 0
    for batch_X, batch_y in train_dataloader: 
        batch_X, batch_y = batch_X.to(device), batch_y.to(device)       

        for x_t, y_t in zip(batch_X, batch_y):

            # number of context pairs
            Nc = random.randint(1,Npts)

            # create random subset of x_t and y_t
            x_c, y_c = zip(*random.sample(list(zip(x_t,y_t)), Nc))
            x_c = torch.stack(x_c)
            x_c = x_c.view(-1,torch.numel(x_c))
            y_c = torch.stack(y_c)
            y_c = y_c.view(-1,torch.numel(y_c))

            for x, y in zip(x_t, y_t):

                x = x.view(-1,torch.numel(x))
                y = y.view(-1,torch.numel(y))

                # compute outputs of model
                output = model(x_c, y_c, x)

                # compute training loss
                train_loss = criterion(output, y)
                
                # reset the gradients back to zero
                # PyTorch accumulates gradients on subsequent backward passes
                optimizer.zero_grad()

                # compute accumulated gradients
                train_loss.backward()
                
                # perform parameter update based on current gradients
                optimizer.step()
                
                # add the mini-batch training loss to epoch loss
                loss += train_loss.item()

    # compute the epoch training loss
    loss = loss/len(train_dataloader)
    
    # display the epoch training loss
    print("epoch : {}/{}, loss = {:.6f}".format(epoch + 1, MAX_EPOCH, loss))