rewrite memory count without implicitly using dimensions by @ic-synth

Author: bm-synth

src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py

@@ -41,9 +41,7 @@ class CogVideoXSafeConv3d(nn.Conv3d):
     """
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
-        memory_count = (
-            (input.shape[0] * input.shape[1] * input.shape[2] * input.shape[3] * input.shape[4]) * 2 / 1024**3
-        )
+        memory_count = torch.prod(torch.tensor(input.shape)) * 2 / 1024**3
 
         # Set to 2GB, suitable for CuDNN
         if memory_count > 2: