How to change the model outcome dtype from int64 to uint8 or float32

[fschvart]

I'm training a binary semantic segmentation model using the PIDnet architecture. I'm not sure why, but even though I only have 0s and 1s, my model output is still int64. This is a problem for me because it's causing some problems when I try to convert the model to TensorRT. Is there a way to change the dtype of the model output?

Thanks!