import torch


def apply_mask(mem_y, out, n_classes):
	"""
	Apply mask on the predicted outputs based on the given task - assuming task-incremental learning setup
	:param mem_y: Actual labels
	:param out: Predictions
	:param n_classes: Number of classes per task
	:return: Masked predictions
	"""
	for i,y in enumerate(mem_y):
		mask = torch.zeros_like(out[i])
		mask[y-(y%n_classes):y-(y%n_classes)+n_classes]=1
		out[i] = out[i].masked_fill(mask == 0, -1e10)
	return out