Revert Sampling in Bernoulli model

successar · Aug 15, 2020 · 894bfba · 894bfba
1 parent 7573197
commit 894bfba
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/Rationale_model/models/classifiers/enc_gen_model.py b/Rationale_model/models/classifiers/enc_gen_model.py
@@ -57,7 +57,7 @@ def forward(self, document, kept_tokens, rationale=None, label=None, metadata=No
         prob_z = kept_tokens.float() + prob_z * (1 - kept_tokens)
         sampler = D.bernoulli.Bernoulli(probs=prob_z)
 
-        sample_z = mask.float()
+        sample_z = sampler.sample() * mask.float()
         encoder_dict = self._encoder(sample_z=sample_z, label=label, metadata=metadata)
 
         loss = 0.0