Probably I misunderstood your use case. I thought you wanted to sample using the probabilities.
In case you just want to keep the two channels with the highest probabilities, you could use a mask or slice the input. This depends on your current use case.
batch_size = 1
c, h, w = 4, 5, 5
x = torch.randn(batch_size, c, h, w)

probs= torch.tensor([0.5, 0.3, 0.1, 0.1])
_, idx = torch.topk(probs, 2)

# Slicing
modelB = nn.Conv2d(2, 1, 3, 1, 1, bias=False)
weight = modelB.weight.clone()
outputB = modelB(x[:, idx])

# Using a mask
modelA = nn.Conv2d(c, 1, 3, 1, 1, bias=False)

# Set weights to "valid" channels to compare the outputs
with torch.no_grad():
    modelA.weight[:, idx] = weight

mask = torch.zeros(batch_size, c, 1, 1)
mask[:, idx] = 1.
outputA = modelA(x * mask)

torch.allclose(outputA, outputB)
> True