Merge pull request #7 from aock/fix-memory

PoseTracker.get_memory() fix

src/loc_ndf/utils/registration.py

@@ -171,13 +171,14 @@ def registration_step(self, points: torch.Tensor, GM_k=None):
 
     def get_memory(self):
         param_size = 0
+        buffer_size = 0
 
         for model in self.models:
             for param in model.parameters():
                 param_size += param.nelement() * param.element_size()
-            buffer_size = 0
             for buffer in model.buffers():
                 buffer_size += buffer.nelement() * buffer.element_size()
+
         size_all_mb = (param_size + buffer_size) / 1024**2
         return size_all_mb
 ##############################################################################