2d sampling

flatironinstitute · Dec 16, 2024 · 0a21fe4 · 0a21fe4
1 parent 9898e73
commit 0a21fe4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/cryo_sbi/utils/estimator_utils.py b/src/cryo_sbi/utils/estimator_utils.py
@@ -87,7 +87,7 @@ def sample_posterior(
         samples = estimator.sample(
             image_batch.to(device, non_blocking=True), shape=(num_samples,)
         ).cpu()
-        theta_samples.append(samples.reshape(-1, image_batch.shape[0]))
+        theta_samples.append(samples)
 
     return torch.cat(theta_samples, dim=1)