Thank you for the reply!
Sorry I should have provided the code snippet in my question but the below are the snippet:
import gc
import math
import os

import numpy as np
import pandas as pd
import torch
from torch.utils.data.dataset import Dataset

class Autoencoder(torch.nn.Module):
  def __init__(self):
    super(Autoencoder, self).__init__()
    self.num_layers = 8 
    self.factor = 0.8
    self.drop_out_percent = 0.2
    self.input_size = 100
    self.activation_layer = torch.nn.Mish()
    self.encoder = torch.nn.Sequential()
    in_size = self.input_size
    for l in range(self.num_layers):
      self.encoder.add_module(f"encoder_block{l}", torch.nn.Sequential(torch.nn.Linear(in_size, round(in_size * self.factor)),
                                                                      self.activation_layer,
                                                                      torch.nn.LayerNorm(round(in_size * self.factor)),
                                                                      torch.nn.Dropout(self.drop_out_percent)))
      in_size = round(in_size*self.factor)
    self.encoder.add_module(f"encoder_block_last", torch.nn.Sequential(torch.nn.Linear(in_size, 2000),
                                                                       self.activation_layer,
                                                                       torch.nn.LayerNorm(2000),
                                                                       torch.nn.Dropout(self.drop_out_percent)))

    self.final_size = in_size

  def forward(self, x):
    return self.encoder(x)

class Output_Layers(torch.nn.Module):
  def __init__(self):
    super(Output_Layers, self).__init__()
    self.num_layers = 10
    self.input_size = 2000
    self.factor = math.exp(np.log(1 / self.input_size) / self.num_layers)
    self.activation_layer = torch.nn.Mish()
    self.decoder = torch.nn.Sequential()
    in_size = self.input_size
    for l in range(self.num_layers):
      if l != (self.num_layers - 1):
        self.decoder.add_module(f"decoder_block{l}", torch.nn.Sequential(torch.nn.Linear(in_size, round(in_size * self.factor)),
                                                                        self.activation_layer,
                                                                        torch.nn.BatchNorm1d(round(in_size * self.factor))))
        in_size = round(in_size*self.factor)
      else:
        self.decoder.add_module("output", torch.nn.Linear(in_size, 1))

  def forward(self, x):
    return self.decoder(x)

class Driver_Model(torch.nn.Module):
  def __init__(self,  autoencoder, output_layer):
    super(Driver_Model, self).__init__()
    self.autoencoder = autoencoder
    self.output_layer = output_layer

  def forward(self, x):
    output = self.autoencoder(x).contiguous()
    output = self.output_layer(output).squeeze()
    return output

class My_Dataset(Dataset):
  def __init__(self, X):
    self.X = X
    
  def __len__(self):
    return(len(self.X))

  def __getitem__(self, idx):
    return self.X[idx]

os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
data_length = 10000
data = []
for i in range(data_length):
  data.append(torch.rand(100))

autoencoder = Autoencoder()
output = Output_Layers()

model_parameter_list = list(autoencoder.parameters()) + list(output.parameters())

model = Driver_Model(autoencoder, output)
model.cuda().double()

test_dataset = My_Dataset(data)
test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size = 64, shuffle = False, num_workers = 4, pin_memory = True, drop_last = True)
print(model)

for j, data in enumerate(test_dataloader):
  data = data.cuda().double()
  output = model(data)

Many thanks!