How to categorize tensor float values into long values?

1

I am working on a multiclass segmentation problem. My target tensors are of the shape ([1,256,256]), with the unique values ([0.0000, 0.0039, 0.0078]). However, I want to convert the float values into long values as ([0, 1, 2]). When I use .long(), all values turn to zeros.
below I have an example of my tensors:

([[0.0039, 0.0039, 0.0039, 0.0000, 0.0000],
[0.0039, 0.0039, 0.0039, 0.0000, 0.0000],
[0.0039, 0.0039, 0.0039, 0.0000, 0.0000],
[0.0000, 0.0078, 0.0078, 0.0078, 0.0000],
[0.0000, 0.0078, 0.0078, 0.0078, 0.0000]])

The desired output should be like the following:

([[1, 1, 1, 0, 0],
[1, 1, 1, 0, 0],
[1, 1, 1, 0, 0],
[0, 2, 2, 2, 0],
[0, 2, 2, 2, 0]])

2

That is expected behavior for .long(), you may want to do something more complicated like the following:

import torch

def convert_float_to_labels(x):
    # Create a mapping tensor for the exact float values
    mapping = torch.tensor([0.0000, 0.0039, 0.0078])
    
    # Find the closest value for each element
    # This handles floating point precision issues
    distances = torch.abs(x.unsqueeze(-1) - mapping)
    indices = torch.argmin(distances, dim=-1)
    
    return indices

# Example usage
x = torch.tensor([
    [0.0039, 0.0039, 0.0039, 0.0000, 0.0000],
    [0.0039, 0.0039, 0.0039, 0.0000, 0.0000],
    [0.0039, 0.0039, 0.0039, 0.0000, 0.0000],
    [0.0000, 0.0078, 0.0078, 0.0078, 0.0000],
    [0.0000, 0.0078, 0.0078, 0.0078, 0.0000]
])

result = convert_float_to_labels(x)
print(result)