Hi @ptrblck,
Thanks a lot for the quick response. Here is my full code:
import torch
from torch import nn
import csv
import copy

class Net(nn.Module):
    
    def __init__(self,param_path,state_dict_map_path):
        
        # run nn.Module's constructor
        
        super(Net,self).__init__()
        
        # path to .pkl file
        
        self.param_path = param_path
        
        # path to .csv file
        
        self.state_dict_map_path = state_dict_map_path
        
        # build net
        
        in_channels = 16
        
        conv1 = nn.Conv2d(in_channels = 1,
                         out_channels = in_channels,
                         kernel_size = (3,3),
                         padding = (1,1))
        
        batch_norm = nn.BatchNorm2d(num_features = 16)
        
        activation = nn.ReLU()
        
        conv2 = nn.Conv2d(in_channels = in_channels,
                         out_channels = in_channels,
                         kernel_size = (3,3),
                         padding = (1,1))
        
        pooling = nn.MaxPool2d(kernel_size = (2,2))
        
        # first stage
        
        stages = [nn.Sequential(conv1,
                                batch_norm,
                                activation,
                                conv2,
                                batch_norm,
                                activation,
                                pooling)]
        
        # next 4 stages
        
        for i in range(4):
            
            conv1 = nn.Conv2d(in_channels = in_channels,
                              out_channels = in_channels * 2,
                              kernel_size = (3,3),
                              padding = (1,1))
            
            batch_norm = nn.BatchNorm2d(num_features = in_channels * 2)
            
            conv2 = nn.Conv2d(in_channels = in_channels * 2,
                              out_channels = in_channels * 2,
                              kernel_size = (3,3),
                              padding = (1,1))
            
            stages += [nn.Sequential(conv1,
                                     batch_norm,
                                     activation,
                                     conv2,
                                     batch_norm,
                                     activation,
                                     pooling)]
            
            in_channels = in_channels * 2
            
        # 6th stage, in_channels = 256
        
        conv1 = nn.Conv2d(in_channels = in_channels,
                          out_channels = in_channels * 2,
                          kernel_size = (3,3),
                          padding = (1,1))
            
        batch_norm = nn.BatchNorm2d(num_features = in_channels * 2)
            
        stages += [nn.Sequential(conv1,
                                 batch_norm,
                                 activation,
                                 pooling)]
        
        in_channels = in_channels * 2
        
        # final stage, in_channels = 512
        
        conv1 = nn.Conv2d(in_channels = in_channels,
                          out_channels = in_channels * 2,
                          kernel_size = (2,2),
                          padding = (0,0))
            
        batch_norm = nn.BatchNorm2d(num_features = in_channels * 2)
            
        stages += [nn.Sequential(conv1,
                                 batch_norm,
                                 activation)]
        
        # assign names to the stages for the state_dict
        
        self.stage1 = stages[0]
        self.stage2 = stages[1]
        self.stage3 = stages[2]
        self.stage4 = stages[3]
        self.stage5 = stages[4]
        self.stage6 = stages[5]
        self.stage7 = stages[6]
        
        # load the pre-trained parameters
        
        self._load_parameters()
    
    # assign parameters from file to layers
    
    def _load_parameters(self):
        
        self.old_state_dict = torch.load(f = self.param_path,
                                         map_location = torch.device('cpu'))
        
        # make a copy to use load_state_dict() method later
        
        state_dict = copy.deepcopy(self.state_dict())
        
        with open(self.state_dict_map_path,newline='') as csvfile:
            reader = csv.reader(csvfile,delimiter=',')
            for old_key,new_key in reader:
                state_dict[new_key] = self.old_state_dict[old_key]
        
        self.load_state_dict(state_dict)
    
    def forward(self,x):
        
        x = self.stage1(x)
        x = self.stage2(x)
        x = self.stage3(x)
        x = self.stage4(x)
        x = self.stage5(x)
        x = self.stage6(x)
        x = self.stage7(x)
        
        # average of all pixels in each feature map
        
        x = nn.functional.avg_pool2d(input = x,
                                     kernel_size = x.shape[2:])
        
        # flatten from N x 1024 x 1 x 1 to N x 1024
        
        x = torch.flatten(input = x,
                           start_dim = 1)
        
        return x

net = Net(param_path,state_dict_map_path)

My code is very similar to yours, except I included the self._load_parameters() line at the end of the constructor instead of running this method outside the class. Does this make a difference? If not, any other suggestions?