From 41fb08f1ccbecfa3ba87b8b3cb82340a49c45421 Mon Sep 17 00:00:00 2001 From: Matthias Fabian Meyer-Bender Date: Mon, 18 Nov 2024 08:49:11 +0000 Subject: [PATCH] Made PyTorch version check compatible with post-release versions --- cellpose/dynamics.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cellpose/dynamics.py b/cellpose/dynamics.py index 8775b919..3132757e 100644 --- a/cellpose/dynamics.py +++ b/cellpose/dynamics.py @@ -620,7 +620,8 @@ def remove_bad_flow_masks(masks, flows, threshold=0.4, device=torch.device("cpu" device0 = device if masks.size > 10000 * 10000 and (device is not None and device.type == "cuda"): - major_version, minor_version, _ = torch.__version__.split(".") + torch_version = torch.__version__.split(".") + major_version, minor_version = torch_version[0], torch_version[1] torch.cuda.empty_cache() if major_version == "1" and int(minor_version) < 10: # for PyTorch version lower than 1.10