It is a bit tricky, requiring shape manipulations
x shape should be like [batch_dims, ngroups, 1, group_size_in]
weight shape at init: [ngroups, group_size_out, group_size_in]
weight shape in forward(): [ngroups, group_size_in, group_size_out]
this thing about weights allows to do init like:
nn.init.kaiming_(weight.view(-1,group_size_in))
then matmul internally reshapes this into bmm format, does bmm and finally outputs [batch_dims,ngroups,1,group_size_out]
PS group conv may actually be ok for a small group count, without this mess, not sure