Hi,
I replaced the default implementation of batchNorm by  this one and I added to it other functions that store the running_mean and running_var during the training however before I was using the Tracing to convert the module, but with the new implementation, it does not work viewing that Tracing does not record functions which are data-dependent or any control-flow, I tried to use Scripting instead but I got the error below.
The function that saves the running_mean and running variance which I call inside the batchNorm:
@torch.jit.script
def save_statisctics(x : torch.Tensor):
    random_string = ''.join(random.choice(string.ascii_lowercase) for i in range(16))
    current_directory = os.path.dirname(os.path.realpath(__file__))

    final_directory = os.path.join(current_directory, r'batchStatistics')
    if not os.path.exists(final_directory):
        os.makedirs(final_directory)

    save_x = x.numpy()
    path = os.path.join(final_directory, '{:s}'.format(random_string))

    np.save(path, save_x)

Error: torch.jit.frontend.UnsupportedNodeError: GeneratorExp aren’t supported:
Please what are the changes that I should make in both BatchNorm script and the function “save_statisctics” in order to make it work? And is there another way than Scripting?
Thank you in advance!