Solved by harsha_g in post #2 
              
              
                Please notice how an object derived from the nn.Module base class is instantiated and called. 
import torch
import torch.nn as nn
 
from heinsen_routing import Routing

class GECorrection(nn.Module):

    def __init__(self, d_cov, d_inp, d_out, n_out):
        super().__init__()
        
        sel…