could anyone please share simple code snippet of callback from using model weights aggregation for torch trainer? Thanks!