From 20a99771bfddce9a9cd08b4a49ee04e2f8ec70ed Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Fri, 11 Apr 2025 19:33:58 -0400 Subject: [PATCH] Add back in padding --- library/custom_train_functions.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/library/custom_train_functions.py b/library/custom_train_functions.py index 40ba51df..43871c84 100644 --- a/library/custom_train_functions.py +++ b/library/custom_train_functions.py @@ -804,6 +804,9 @@ class WaveletLoss(nn.Module): # Combine high frequency bands for visualization if combined_hf_pred and combined_hf_target: + combined_hf_pred = self._pad_tensors(combined_hf_pred) + combined_hf_target = self._pad_tensors(combined_hf_target) + combined_hf_pred = torch.cat(combined_hf_pred, dim=1) combined_hf_target = torch.cat(combined_hf_target, dim=1) else: