minor_changes_compress_ar

spherical_models/compression_models/sphere_scale_hyperprior.py

@@ -271,6 +271,10 @@ class SphereScaleHyperprior(SphereCompressionModel):
         assert len(self.autoregressive.list_conv) == 2, "only 2-hop convolution supported"
         neighbors_indices = neighbors_indices.to(y_hat.device)
         neighbors_weights = neighbors_weights.to(y_hat.device)
+        
+        cdf = self.gaussian_conditional._quantized_cdf.tolist()
+        cdf_lengths = self.gaussian_conditional._cdf_length.tolist()
+        offsets = self.gaussian_conditional._offset.tolist()
 
         encoder = BufferedRansEncoder()
         symbols_list = []
@@ -311,12 +315,9 @@ class SphereScaleHyperprior(SphereCompressionModel):
             y_q = self.gaussian_conditional.quantize(y_hat[:, n, :], "symbols") # TODO if means: , means_hat)
             y_hat[:, n, :] = y_q # TODO if means: + means_hat
 
-            symbols_list.extend(y_q.reshape(-1).int().tolist())
-            indexes_list.extend(indexes.reshape(-1).int().tolist())
+            symbols_list.extend(y_q.squeeze().tolist())
+            indexes_list.extend(indexes.squeeze().tolist())
 
-        cdf = self.gaussian_conditional._quantized_cdf.tolist()
-        cdf_lengths = self.gaussian_conditional._cdf_length.reshape(-1).int().tolist()
-        offsets = self.gaussian_conditional._offset.reshape(-1).int().tolist()
         encoder.encode_with_indexes(
             symbols_list,
             indexes_list,