Default diff-sampling to gumbel softmax

This make it consistent with the implementation used in https://proceedings.mlr.press/v181/lang22a/lang22a.pdf See also: #5
braun-steven · Nov 8, 2023 · 29680f3 · 29680f3
1 parent 179cc72
commit 29680f3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/simple_einet/sampling_utils.py b/simple_einet/sampling_utils.py
@@ -204,7 +204,7 @@ def sample_categorical_differentiably(
     tau: float,
     logits: torch.Tensor = None,
     log_weights: torch.Tensor = None,
-    method=DiffSampleMethod.SIMPLE,
+    method=DiffSampleMethod.GUMBEL,
 ) -> torch.Tensor:
     """
     Perform differentiable sampling/mpe on the given input along a specific dimension.