Hi, after the last epoch, i have an array of local model weights, which i use to load them onto an empty model to test it later on.

```
global_model = CNNMnist(args=args)
local_model1 = CNNMnist(args=args)
.........
#training
local_weights, local_losses = [], []
local_model = LocalUpdate(args=args, dataset=train_dataset,
idxs=user_groups[idx], logger=logger)
w, loss = local_model.update_weights(
model=copy.deepcopy(global_model), global_round=epoch)
local_weights.append(copy.deepcopy(w))
local_losses.append(copy.deepcopy(loss))
# update global weights
global_weights = average_weights(local_weights)
# update global weights
global_model.load_state_dict(global_weights)
```

The following line is where i start tinkering:

```
local_model1 = local_model1.load_state_dict(local_weights[0])
```

then i send the model to a function to get the metrics:

`test_acc, test_loss, precision, recall, f1, confusion = sam_eval(args,local_model1,test_dataset)`

the error is thrown right towards the beginning of the function:

```
def sam_eval(args, model, test_dataset):
device = 'cuda' if args.gpu else 'cpu'
criterion = nn.NLLLoss().to(device)
model.eval() #this is the line throwing the error
loss, total, correct = 0.0, 0.0, 0.0
..........
```

p.s. when i do the same thing with `global_model`

, it works just fine. And i also tried adding `local_model.to(device)`

line to my code.

Any help would be really appreciated.