From 045d3fba90a02a9337cad99799c43e7dd084317c Mon Sep 17 00:00:00 2001 From: Taeyup Song Date: Wed, 5 Mar 2025 22:14:04 +0900 Subject: [PATCH] ADD exception handler for efficientnet model --- src/main.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/src/main.py b/src/main.py index c2265ec..d6f004d 100755 --- a/src/main.py +++ b/src/main.py @@ -212,12 +212,11 @@ else: for key in v.keys(): v[key] = update_avg[key] + v[key] * args.server_momentum - #new_weights = deepcopy(model.state_dict()) - #for key in new_weights.keys(): - # new_weights[key] = new_weights[key] - v[key] * args.server_lr - #model.load_state_dict(new_weights) - for key in model.state_dict(): - model.state_dict()[key] -= v[key] * args.server_lr + + for key in model.state_dict(): + last_dot_index = key.rfind('.') + if key[last_dot_index + 1:] != "num_batches_tracked": + model.state_dict()[key] -= v[key] * args.server_lr # Compute round average loss and accuracies if round % args.server_stats_every == 0: