Solved by user_123454321 in post #2 
              
              
                Can you replace 
            model.log_weights(step, writer)

with 
            model.module.log_weights(step, writer)