code: def apply_threshold(preds, t):
      assert np.allclose(preds.sum(-1), np.ones(preds.shape[0]))
      prob = preds[:, -1]
      thresholded = (prob >= t).astype(int)
      preds = np.zeros_like(preds)
      preds[np.arange(len(thresholded)), thresholded.reshape(-1)] = 1
      return preds
  
file_path: megatron-lm/tasks/ensemble_classifier.py
imports: import os
  import argparse
  import collections
  import numpy
  import torch
label_desc: Use image node_id 4 for calling a global function with example usage: apply_threshold(preds, t) and returns: preds