Bug fixes

1. Turn off autograd while decoding canonical and pose features
2. Change default batch size to (4, 8)

2 files changed, 11 insertions, 10 deletions

diff --git a/config.py b/config.py
index c93fe19..991a3a6 100644
--- a/config.py
+++ b/config.py
@@ -33,7 +33,7 @@ config: Configuration = {
         # Batch size (pr, k)
         # `pr` denotes number of persons
         # `k` denotes number of sequences per person
-        'batch_size': (2, 2),
+        'batch_size': (4, 8),
         # Number of workers of Dataloader
         'num_workers': 4,
         # Faster data transfer from RAM to GPU if enabled
diff --git a/models/auto_encoder.py b/models/auto_encoder.py
index 5e7558b..36be868 100644
--- a/models/auto_encoder.py
+++ b/models/auto_encoder.py
@@ -132,15 +132,16 @@ class AutoEncoder(nn.Module):
         # x_c1_t2 is the frame for later module
         (f_a_c1_t2, f_c_c1_t2, f_p_c1_t2) = self.encoder(x_c1_t2)
 
-        # Decode canonical features for HPM
-        x_c_c1_t2 = self.decoder(
-            torch.zeros_like(f_a_c1_t2), f_c_c1_t2, torch.zeros_like(f_p_c1_t2),
-            no_trans_conv=True
-        )
-        # Decode pose features for Part Net
-        x_p_c1_t2 = self.decoder(
-            torch.zeros_like(f_a_c1_t2), torch.zeros_like(f_c_c1_t2), f_p_c1_t2
-        )
+        with torch.no_grad():
+            # Decode canonical features for HPM
+            x_c_c1_t2 = self.decoder(
+                torch.zeros_like(f_a_c1_t2), f_c_c1_t2, torch.zeros_like(f_p_c1_t2),
+                no_trans_conv=True
+            )
+            # Decode pose features for Part Net
+            x_p_c1_t2 = self.decoder(
+                torch.zeros_like(f_a_c1_t2), torch.zeros_like(f_c_c1_t2), f_p_c1_t2
+            )
 
         if self.training:
             # t1 is random time step, c2 is another condition