Both approaches will work the same, if you are using a pure functional approach.
A custom nn.Module class is useful in case your custom module needs to store parameters, buffers etc.
E.g. nn.Conv2d itself is defined as an nn.Module, since it stores the weight, bias parameters as well as some internal attributes, e.g. padding, stride, etc.
In your current example you are just applying some modules, so you can easily define it as a function.