From e03564e426b825a4174f1a943c8ada4e75593416 Mon Sep 17 00:00:00 2001 From: um3 Date: Wed, 24 Jan 2024 17:13:11 +0800 Subject: [PATCH] rm cpu --- utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/utils.py b/utils.py index ce5bcab..885f402 100644 --- a/utils.py +++ b/utils.py @@ -56,10 +56,10 @@ def save_network(network, dirname, epoch_label, local_rank=-1): if local_rank>-1: if local_rank == 0: # save the main process model - torch.save(network.cpu().state_dict(), save_path) + torch.save(network.state_dict(), save_path) network.cuda(local_rank) else: - torch.save(network.cpu().state_dict(), save_path) + torch.save(network.state_dict(), save_path) network.cuda()