def forward(self, dtype):
    logit_scale = (self.max_log_temp - F.softplus((self.max_log_temp - self.logit_scale)))
    return logit_scale.exp().type(dtype)
def forward(self, dtype):
    logit_scale = (self.max_log_temp - F.softplus((self.max_log_temp - self.logit_scale)))
    return logit_scale.exp().type(dtype)