So both model setup functions haven’t been working because of torch.legacy.nn, but they never gave any errors because for some reason this if statement isn’t working properly:
if isinstance(layer, nn.ReLU):