How can I apply group normalization after a full-connection layer?

Your code is correct, but let's see what happens in a small example.

The output of a fully-connected layer is usually a 2D-tensor with shape (batch_size, hidden_size) so I will focus on this kind of input, but remember that GroupNorm supports tensors with an arbitrary number of dimensions. In fact, GroupNorm works always on the last dimension of the tensor.

GroupNorm treats all the samples in the batch as independent and it creates n_groups from the last dimension of the tensor, as you can see from the image.

When the input tensor is 2D, the cube in the image becomes a square because there is no third vertical dimension, so in practice the normalization is performed on fixed-size consecutive pieces of the rows of the input matrix.

Let's see an example with some code.

import torch
import torch.nn as nn

batch_size = 2
hidden_size = 32
n_groups = 8
group_size = hidden_size // n_groups # = 4

# Input tensor that can be the result of a fully-connected layer
x = torch.rand(batch_size, hidden_size)

# GroupNorm with affine disabled to simplify the inspection of results
gn1 = nn.GroupNorm(n_groups, hidden_size, affine=False)
r = gn1(x)

# The rows are split into n_groups (8) groups of size group_size (4)
# and the normalization is applied to these pieces of rows.
# We can check it for the first group x[0, :group_size] with the following code
first_group = x[0, :group_size]
normalized_first_group = (first_group - first_group.mean())/torch.sqrt(first_group.var(unbiased=False) + gn1.eps)
print(r[0, :4])
print(normalized_first_group)
if(torch.allclose(r[0, :4], normalized_first_group)):
    print('The result on the first group is the expected one')