In case it’s useful for future readers, here’s a generalized function inspired by the solution from @pascal_notsawo. In my-use case, I needed the 3D indices of each sample’s maximum value (as opposed to masking them in the original tensor like the OP requested). I also handled the edge case of multiple, identical maximum maximum values
from collections import defaultdict
import torch


def sample_max_idxs(batch, on_duplicate='first'):
    """Get indices of maximum sample values

    A combination of operations inspired from the following posts
    https://discuss.pytorch.org/t/masking-batch-wise-max/108291/2
    https://stackoverflow.com/a/50046640/7446465

    :param torch.Tensor[N, ...] batch: Batch of ND tensor samples
    :param str on_duplicate: how to handle exact tie in maximum values. Valid options are
        * first: keep the lowest 2D index
        * last: keep the highest 2D index
        * None: don't remove duplicates
    :return torch.Tensor[N, len(batch.shape)] idxs: Coordinates for the max value within
        each sample of the batch
    """

    # Compress all secondary axes
    batch_size, *other_axes = batch.shape
    batch = batch.view(batch_size, -1)

    # Create a 1D mask for the max indices
    max_per_sample = batch.max(dim=1, keepdim=True)[0]
    mask = torch.zeros_like(batch)
    mask[batch == max_per_sample] = 1

    # Reshape to original dims and use ``nonzero`` binary trick to generate ND indices
    mask = mask.view(batch_size, *other_axes)
    idxs = torch.nonzero(mask * mask)

    # Handle duplicates
    batch_idxs = torch.sort(idxs[:, 0])[0]
    if on_duplicate is None or batch_idxs.numel() == batch_size:
        return idxs
    valid_duplicates = ['first', 'last']
    if on_duplicate not in valid_duplicates:
        raise ValueError(f'Expected on_duplicate to be one of {",".join(valid_duplicates)}, got {on_duplicate}')
    batch2locs = defaultdict(list)
    offset = 0
    for i, batch_idx in enumerate(batch_idxs):
        if i + offset != batch_idx:
            offset -= 1
        batch2locs[batch_idx.item()].append(i)
    remove = []
    for locs in batch2locs.values():
        if len(locs) > 1:
            keep_idx = 0 if on_duplicate == 'first' else len(locs) - 1
            remove.extend([loc for i, loc in enumerate(locs) if i != keep_idx])
    idxs = [sample_idx for i, sample_idx in enumerate(idxs) if i not in remove]
    return torch.stack(idxs)