Current code only checks image tensor for CUDA device If label tensor is on GPU (even if image is on CPU), torch.cuda.empty_cache() is not called This can cause GPU memory leaks during connected component analysis (in auto3dseg/analyzer.py)