def forward(self, x):
    x = self.mlp(x)
    x = nn.functional.normalize(x, dim=(- 1), p=2)
    x = self.last_layer(x)
    return x
def forward(self, x):
    x = self.mlp(x)
    x = nn.functional.normalize(x, dim=(- 1), p=2)
    x = self.last_layer(x)
    return x