Fixbug computing Kurtosis

lvapeab · Jun 17, 2020 · 6d20321 · 6d20321
1 parent 7c8ddd5
commit 6d20321
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/active_learning_nmt.py b/active_learning_nmt.py
@@ -113,7 +113,7 @@ def sampling_function(source_stream, n_samples, mode='random', src_seqs=None, hy
     elif mode == 'attention-distraction':
         assert alphas_stream is not None, 'Cannot use sampling mode "%s" without alphas.' % str(mode)
         # Get the mean minus curtosis for each sentence
-        stream_scores = [-kurtosis(alphas).mean() for alphas in alphas_stream]
+        stream_scores = [-kurtosis(alphas, axis=-1).mean() for alphas in alphas_stream]
         # Return the maximum values
         indices_to_return = np.asarray(stream_scores).argsort()[::-1][:n_samples]  # Get the maximum values
         return indices_to_return, [None]
@@ -164,7 +164,7 @@ def sampling_function(source_stream, n_samples, mode='random', src_seqs=None, hy
             stream_scores_cp[i] = float(cp_penalty) / len(src_seq)
 
         # Get the mean minus curtosis for each sentence
-        stream_scores_ad = [-kurtosis(alphas).mean() for alphas in alphas_stream]
+        stream_scores_ad = [-kurtosis(alphas, axis=-1).mean() for alphas in alphas_stream]
 
         # Random sampling
         indices_to_return_random = np.arange(len(source_stream))