VAE: Support different latent_var_log options when returning intermediate features for 3D perceptual loss

xora/models/autoencoders/video_autoencoder.py

@@ -370,9 +370,7 @@ class Encoder(nn.Module):
                 raise ValueError(f"Invalid input shape: {sample.shape}")
 
         if return_features:
-            features.append(
-                sample[:, sample.shape[1] // 2, ...]
-            )  # Add the latent means as final feature
+            features.append(sample[:, : self.latent_channels, ...])
             return sample, features
         return sample