From efe1246323be451155a464fff84787d3d778f387 Mon Sep 17 00:00:00 2001
From: Super Daniel <78588128+super-dainiu@users.noreply.github.com>
Date: Fri, 10 Feb 2023 10:49:44 +0800
Subject: [PATCH] Fix issue #5299. (#5343)

---
 nni/compression/pytorch/speedup/infer_mask.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/nni/compression/pytorch/speedup/infer_mask.py b/nni/compression/pytorch/speedup/infer_mask.py
index ead1ffb73..57c01d752 100644
--- a/nni/compression/pytorch/speedup/infer_mask.py
+++ b/nni/compression/pytorch/speedup/infer_mask.py
@@ -86,11 +86,15 @@ class AutoMaskInference:
             self.output_mask = self.in_masks[0]
         else:
             if isinstance(self.output, torch.Tensor):
+                if self.output.requires_grad:
+                    self.output.retain_grad()   # issue #5299
                 self.output_mask = torch.ones_like(self.output)
             elif isinstance(self.output, list) or isinstance(self.output, tuple):
                 self.output_mask = []
                 for o_tensor in self.output:
                     if isinstance(o_tensor, torch.Tensor):
+                        if o_tensor.requires_grad:
+                            o_tensor.retain_grad()  # issue #5299
                         self.output_mask.append(torch.ones_like(o_tensor))
                     else:
                         # if one of the outputs is not tensor, set the corresponding