def _get_nearest_neigbor_distr(self, x: Tensor) -> Tensor:
        """
        Find closest prototype for output samples in batch x.
        :param x: Batch of network logits.
        :return: one-hot representation of the predicted class.
        """
        ns = x.size(0)
        nd = x.view(ns, -1).shape[-1]

        # Get prototypes
        seen_c = len(self.p_mem.keys())
        if seen_c == 0:  # no prototypes yet, output uniform distr. all classes
            return (
                torch.Tensor(ns, self.n_classes)
                .fill_(1.0 / self.n_classes)
                .to(x.device)
            )
        means = torch.ones(seen_c, nd).to(x.device) * float("inf")
        for c, c_proto in self.p_mem.items():
            means[c] = c_proto  # Class idx gets allocated its prototype

        # Predict nearest mean
        classpred = torch.LongTensor(ns)
        for s_idx in range(ns):  # Per sample
            dist = -torch.mm(means, x[s_idx].unsqueeze(-1))  # Dot product
            _, ii = dist.min(0)  # Min dist (no proto = inf)
            ii = ii.squeeze()
            classpred[s_idx] = ii.item()  # Allocate class idx

        # Convert to 1-hot
        out = torch.zeros(ns, self.n_classes).to(x.device)
        for s_idx in range(ns):
            out[s_idx, classpred[s_idx]] = 1
        return out  # return 1-of-C code, ns x nc

        Find closest prototype for output samples in batch x.
        :param x: Batch of network logits.
        :return: one-hot representation of the predicted class.
        
avalanche/training/plugins/cope.py