If your metric is stateless (i.e. you don’t want to store arguments or other parameters), then you could just define a method. Otherwise, a custom class derived from nn.Module could store some arguments e.g. reductions as an internal attribute.
Also, have a look at this post I’ve written some time ago where I discuss the advantages of both approaches.
1 Like