fix: 🐛 pad to concatenate for tensorboard

local.py

@@ -4,6 +4,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
+import math
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 from PIL import Image
 from torch.utils.data import Dataset
@@ -86,9 +87,19 @@ def show_random_dataset_image(dataset):
     print("Image size is %s" % {img[0].shape})
     plt.show()
 
+def pad_to_size(small_tensor, target_size):
+    if small_tensor.size() > target_size:
+        msg = f"Can't pad tensor of size {small_tensor.size()} to tensor of size {target_size}."
+        raise ValueError(msg)
+    if small_tensor.size() == target_size:
+        return small_tensor
+    pad_twoside = []
+    for small_size, large_size in zip(small_tensor.shape, target_size):
+        pad_twoside.append(math.floor((large_size - small_size)/2))
+        pad_twoside.append(math.ceil((large_size-small_size)/2))
+    return torch.nn.functional.pad(small_tensor, pad_twoside[::-1])
 
 def apply_and_show_random_image(f, ds):
-
     # pick random raw image from dataset
     img_tensor = ds[np.random.randint(len(ds))][0]
 

solution.py

@@ -40,6 +40,7 @@
     apply_and_show_random_image,
     plot_receptive_field,
     show_random_dataset_image,
+    pad_to_size
 )
 
 # %% tags=[]
@@ -1099,7 +1100,7 @@ def train(
             # check if we log images in this iteration
             if step % log_image_interval == 0:
                 combined_image = torch.cat(
-                    [x.to("cpu"), y.to("cpu"), prediction.to("cpu").detach()], dim=3
+                    [x.to("cpu"), pad_to_size(y.to("cpu"), x.size()), pad_to_size(prediction.to("cpu").detach(), x.size())], dim=3
                 )
                 tb_logger.add_images(
                     tag="input_target_prediction",