text: 
  import torch
  
  
  class ThresholdModule(torch.nn.Module):
  	def __init__(
  			self, num_thresholds: int = 1, threshold_min: float = 0.0, threshold_max: float = 1.0, threshold_delta: float = 0.01):
  		super().__init__()
  		self.num_thresholds = num_thresholds
  		self.threshold_min = threshold_min
  		self.threshold_max = threshold_max
  		self.threshold_delta = threshold_delta
  
  		self._thresholds = torch.nn.Parameter(
  			torch.zeros(num_thresholds, dtype=torch.float32),
  			requires_grad=False
  		)
  
  	@property
  	def thresholds(self):
  		return self._thresholds
  
  	def forward(self, scores):
  		return self.predict(scores, self.thresholds)
  
  	def predict(self, scores, thresholds):
  		preds = scores.gt(thresholds).long()
  		return preds
  
  	def update_thresholds(self, new_value):
  		self.thresholds.copy_(new_value)
  
  	def get_range(self, threshold_min: float = None, threshold_max: float = None, threshold_delta: float = None):
  		threshold_range = torch.arange(
  			self.threshold_min if threshold_min is None else threshold_min,
  			self.threshold_max if threshold_max is None else threshold_max,
  			self.threshold_delta if threshold_delta is None else threshold_delta,
  			dtype=torch.float32,
  			requires_grad=False
  		)
  		return threshold_range
  
  	def get_range_predictions(self, scores, threshold_min: float = None, threshold_max: float = None, threshold_delta: float = None):
  		if self.num_thresholds == 1:
  			t_preds = self.get_range_threshold_predictions(scores, threshold_min, threshold_max, threshold_delta)
  			for threshold, preds in t_preds:
  				yield threshold, preds
  		else:
  			raise NotImplementedError('Still need to implement this properly')
  			# for t_idx in range(self.num_thresholds):
  			# 	t_scores = scores[..., t_idx]
  			# 	t_range_preds = self.get_range_threshold_predictions(t_scores)
  			# 	for t_threshold, t_preds in t_range_preds:
  			# 		yield t_threshold, t_preds
  
  	def get_range_threshold_predictions(self, scores, threshold_min: float = None, threshold_max: float = None, threshold_delta: float = None):
  		for threshold in self.get_range(threshold_min, threshold_max, threshold_delta):
  			preds = self.predict(scores, threshold)
  			yield threshold, preds
  
  