Fix BatchMemoryManager length

Author: Dariush Wahdany

opacus/utils/batch_memory_manager.py

@@ -16,12 +16,13 @@
 from typing import List
 
 import numpy as np
+from torch.utils.data import BatchSampler, DataLoader, Sampler
+
 from opacus.optimizers import DPOptimizer
 from opacus.utils.uniform_sampler import (
     DistributedUniformWithReplacementSampler,
     UniformWithReplacementSampler,
 )
-from torch.utils.data import BatchSampler, DataLoader, Sampler
 
 
 class BatchSplittingSampler(Sampler[List[int]]):
@@ -71,13 +72,17 @@ def __iter__(self):
     def __len__(self):
         if isinstance(self.sampler, BatchSampler):
             return int(
-                len(self.sampler) * (self.sampler.batch_size / self.max_batch_size)
+                np.ceil(
+                    len(self.sampler) * (self.sampler.batch_size / self.max_batch_size)
+                )
             )
         elif isinstance(self.sampler, UniformWithReplacementSampler) or isinstance(
             self.sampler, DistributedUniformWithReplacementSampler
         ):
             expected_batch_size = self.sampler.sample_rate * self.sampler.num_samples
-            return int(len(self.sampler) * (expected_batch_size / self.max_batch_size))
+            return int(
+                np.ceil(len(self.sampler) * (expected_batch_size / self.max_batch_size))
+            )
 
         return len(self.sampler)