diff --git a/codes/utils/loss_accumulator.py b/codes/utils/loss_accumulator.py
index a8762360..8965cf3d 100644
--- a/codes/utils/loss_accumulator.py
+++ b/codes/utils/loss_accumulator.py
@@ -8,17 +8,22 @@ class LossAccumulator:
 
     def add_loss(self, name, tensor):
         if name not in self.buffers.keys():
-            self.buffers[name] = (0, torch.zeros(self.buffer_sz))
-        i, buf = self.buffers[name]
+            self.buffers[name] = (0, torch.zeros(self.buffer_sz), False)
+        i, buf, filled = self.buffers[name]
         # Can take tensors or just plain python numbers.
         if isinstance(tensor, torch.Tensor):
             buf[i] = tensor.detach().cpu()
         else:
             buf[i] = tensor
-        self.buffers[name] = ((i+1) % self.buffer_sz, buf)
+        filled = i+1 >= self.buffer_sz or filled
+        self.buffers[name] = ((i+1) % self.buffer_sz, buf, filled)
 
     def as_dict(self):
         result = {}
         for k, v in self.buffers.items():
-            result["loss_" + k] = torch.mean(v[1])
+            i, buf, filled = v
+            if filled:
+                result["loss_" + k] = torch.mean(buf)
+            else:
+                result["loss_" + k] = torch.mean(buf[:i])
         return result
\ No newline at end of file