Solved by praveenck06 in post #3 
              
              
                Thanks @ptrblck. 
separating fc layer and adding flatten before it helped. 
class CustomLayer(nn.Module):
    def __init__(self): 
        super().__init__()

        
    def forward(self, input_features):
        input_features = input_features*0.5 # some calculations
        return input_features…