How to create a tensor in a custom python function within define_graph

Describe the question.

How do create new torch tensors and have them go to the correct device. I would like to do things like taking the square of the tensor? I found this example:
https://docs.nvidia.com/deeplearning/dali/archives/dali_1_18_0/user-guide/docs/examples/custom_operations/python_operator.html

def edit_images(image1, image2):
    assert image1.shape == image2.shape
    for i in range(c):
        h, w, c = image1.shape
        perturbation = torch.rand(h, w) 
        new_image1 = torch.zeros(h,w,c)
        new_image2 = torch.zeros(h,w,c)
        new_image1[:, :, i] = image1[:, :, i] * torch.square(perturbation)
        new_image2[:, :, i] = image2[:, :,i]  * torch.square(perturbation)
    return new_image1, new_image2

Check for duplicates

• I have searched the open bugs/issues and have found no duplicates for this bug report