Hi Pytorcher!



 pytorcher:

at the end I should do another dictionary check and set input[0] to be mainOut again


This is exactly right (if I understand your use case).  Thus, something
like:
def forward (input):
    if type (input) is dict:
        some_tensor = input[0]
    else:
        some_tensor = input
    
    # do stuff with some_tensor
    # then do
    main_out = self.batchNormLayer (some_tensor)
    
    if type (input) is dict:
        input[0] = main_out
        return input
    else:
        return main_out


Best.
K. Frank