From 61a1d59ac5fe693320360cb7bf913d87e70a87bf Mon Sep 17 00:00:00 2001 From: PNRxA Date: Sat, 25 Apr 2020 16:30:19 +1000 Subject: [PATCH] numpy to use CPU when using CUDA --- utils/generic_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utils/generic_utils.py b/utils/generic_utils.py index 5d91d74d..435d2b10 100644 --- a/utils/generic_utils.py +++ b/utils/generic_utils.py @@ -157,7 +157,7 @@ def check_update(model, grad_clip, ignore_stopnet=False): grad_norm = torch.nn.utils.clip_grad_norm_([param for name, param in model.named_parameters() if 'stopnet' not in name], grad_clip) else: grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip) - if np.isinf(grad_norm): + if np.isinf(grad_norm.cpu()): print(" | > Gradient is INF !!") skip_flag = True return grad_norm, skip_flag