From 7e37f63e5e77a217a08d9261af49e63ddb51d306 Mon Sep 17 00:00:00 2001 From: Aaron Shi Date: Wed, 10 Jan 2024 21:33:56 +0000 Subject: [PATCH] [Reference Cycle Detector] Ignore FakeTensor in cycle leak detection (#117116) Summary: Skip FakeTensors since these tensors are not actually using GPU memory. Reference Cycle Detector does not need to generate plots for these tensors. Test Plan: CI and internal testing. Differential Revision: D52637209 Pull Request resolved: https://github.com/pytorch/pytorch/pull/117116 Approved by: https://github.com/zdevito, https://github.com/tianfengfrank --- torch/utils/viz/_cycles.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch/utils/viz/_cycles.py b/torch/utils/viz/_cycles.py index 80a7d35cd4c..f17348e401c 100644 --- a/torch/utils/viz/_cycles.py +++ b/torch/utils/viz/_cycles.py @@ -310,7 +310,7 @@ def escape(n): def is_cuda_tensor(obj): - return isinstance(obj, torch.Tensor) and obj.is_cuda + return isinstance(obj, torch.Tensor) and obj.is_cuda and not isinstance(obj, torch._subclasses.FakeTensor) def cuda_allocation_context(): snapshot = torch.cuda.memory._snapshot()