Update loss_utils.py (#1247)

633db6be · Duo Li · GitHub · e948c537 · 633db6be
Unverified Commit 633db6be authored Jan 30, 2023 by Duo Li Committed by GitHub Jan 30, 2023
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

pcdet/utils/loss_utils.py pcdet/utils/loss_utils.py +2 -1

No files found.
--- a/pcdet/utils/loss_utils.py
+++ b/pcdet/utils/loss_utils.py
@@ -148,6 +148,7 @@ class WeightedL1Loss(nn.Module):
            self.code_weights = np.array(code_weights, dtype=np.float32)
            self.code_weights = torch.from_numpy(self.code_weights).cuda()

+    @torch.amp.cuda.custom_fwd(cast_inputs=torch.float16)
    def forward(self, input: torch.Tensor, target: torch.Tensor, weights: torch.Tensor = None):
        """
        Args: