Hi,
You can just apply nn.Conv2d() and set groups=in_channels, and out_channels=NK, it should be what you want.
I assume that N=3 and K=5 here.
input = torch.randn(1,3,10,10)
layer = nn.Conv2d(in_channels=3, out_channels=15, 3, groups=3)
out = layer(input)
out.shape
# output torch.Size([1, 15, 10, 10])