Threshold a tensor into binary values using pytorch (torch.where)

binary_tensor.py

import torch

# Create a tensor with values between 0 and 1
tensor = torch.rand(3, 3)

# Set a threshold value
threshold = 0.5

# Use torch.where to threshold the tensor
binary_tensor = torch.where(tensor > threshold, torch.tensor(1), torch.tensor(0))

# Print the original tensor and the thresholded tensor
print(tensor)
print(binary_tensor)