Fix tqdm import

osrt/layers.py

@@ -295,6 +295,7 @@ class SoftPositionEmbed(nn.Module):
     def forward(self, inputs):
         return inputs + self.dense(torch.tensor(self.grid).cuda()).permute(0, 3, 1, 2) #  from [b, h, w, c] to [b, c, h, w]
     
+    
 ### New transformer implementation of SlotAttention inspired from https://github.com/ThomasMrY/VCT/blob/master/models/visual_concept_tokenizor.py
 class TransformerSlotAttention(nn.Module):
     """

osrt/model.py

@@ -54,14 +54,11 @@ def spatial_broadcast(slots, resolution):
     # `grid` has shape: [batch_size*num_slots, width, height, slot_size].
     return grid
 
-
-# TODO : adapt this model
 class SlotAttentionAutoEncoder(nn.Module):
     """
     Slot Attention as introduced by Locatello et al. but with the AutoEncoder part to extract image embeddings.
 
     Implementation inspired from official repo : https://github.com/google-research/google-research/blob/master/slot_attention/model.py
-
     """
 
     def __init__(self, resolution, num_slots, num_iterations):

train_sa.py

@@ -13,7 +13,7 @@ from osrt.utils.common import mse2psnr
 
 from torch.utils.data import DataLoader
 import torch.nn.functional as F
-import tqdm
+from tqdm import tqdm
 
 def train_step(batch, model, optimizer, device):
     """Perform a single training step."""