How do I pass the number of classes to UDA as a parameter?

[TombacS]

I try to use UDA as a comparison algorithm for my work. However, the network parameters of UDA in this module are written in UDA.py, and the default class number is 10. How can I modify it so that I can pass the number of classes into this module so that the size of the last layer of the network automatically adapts?