Fix PyTorch floor division warning

Signed-off-by: Steven Lang <steven.lang.mz@gmail.com>
braun-steven · Nov 24, 2021 · e63521d · e63521d
1 parent c0f7004
commit e63521d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/einet.py b/einet.py
@@ -352,8 +352,8 @@ def sample(
             # Obtain repetition indices
             ctx.repetition_indices = (ctx.parent_indices % self.config.R).squeeze(1)
             # Shift indices
-            ctx.parent_indices = ctx.parent_indices // self.config.R
-            # ctx.parent_indices = ctx.parent_indices % num_roots
+            # ctx.parent_indices = ctx.parent_indices // self.config.R
+            ctx.parent_indices = torch.div(ctx.parent_indices, self.config.R, rounding_mode="floor")
 
             # Now each sample in `indices` belongs to one repetition, index in `repetition_indices`