diff --git a/omnigibson/utils/processing_utils.py b/omnigibson/utils/processing_utils.py
index c47621317..bd81ceba8 100644
--- a/omnigibson/utils/processing_utils.py
+++ b/omnigibson/utils/processing_utils.py
@@ -109,7 +109,7 @@ def reset(self):
     @property
     def state_size(self):
         # This is the size of the internal buffer plus the current index and fully filled single values
-        return th.prod(self.past_samples.shape) + 2
+        return cb.prod(self.past_samples.shape) + 2
 
     def _dump_state(self):
         # Run super init first