diff --git a/osrt/layers.py b/osrt/layers.py
index deaf9d8419b9e18adc6ffffccabeac5afc687f64..19f9395e53cf9aa782fe77a5d7ef729a083f3874 100644
--- a/osrt/layers.py
+++ b/osrt/layers.py
@@ -244,8 +244,7 @@ class SlotAttention(nn.Module):
 
         k, v = self.to_k(inputs), self.to_v(inputs)
 
-        if slots is None:
-            slots = self.slots_mu + torch.exp(self.slots_log_sigma) * torch.randn(len(inputs), self.num_slots, self.slot_size, device = self.slots_mu.device)
+        slots = self.slots_mu + torch.exp(self.slots_log_sigma) * torch.randn(len(inputs), self.num_slots, self.slot_size, device = self.slots_mu.device)
 
         # Multiple rounds of attention.
         for _ in range(self.iters):