Hi, @albanD
The code below is an example.
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.parameter as Parameter
from torch.autograd import Variable


BS, C = 256, 8
sn = torch.randn(1)[0]
gate_rand = Variable(torch.rand(BS, C).cuda(), requires_grad=True)
noise_rand = Variable(torch.rand(BS, C).cuda(), requires_grad=True)

gate = gate_rand + 1.0e-3
noise = noise_rand + 1.0e-4
Hx = gate + sn*noise
print(gate.requires_grad, noise.requires_grad)

topk, indices = torch.topk(Hx, C//2)
_, neg_indices = torch.topk(-Hx, C//2)
topk1, indices1 = torch.topk(Hx, C//2+1)
thresh_k = topk.min(1, keepdim=True)[0]
thresh_k1 = topk1.min(1, keepdim=True)[0]

Hx_mask = Hx.scatter(1, neg_indices, float('-inf'))

kmatrix = thresh_k.repeat(1, C)
kmatrix1 = thresh_k1.repeat(1, C)
print(kmatrix.size(), kmatrix1.size(), Hx.size(), indices.size())
kth_excluding = kmatrix.scatter(1, indices, kmatrix1)
print('Hx: ', Hx_mask.size())
print('kth_excluding: ', kth_excluding.size())