Forward in Model Definition


I am writing a model in Pytorch. At the last layer, I have a fully connected that gives 3 points (let’s assume {x, y, h} ).

Now, I want to add a function in a forward that it creates a tensor of size 5*5 which has the value of axes in its. (for example tensor[1,1]=[1,1] ). I want to compute the euclidean distance of this tensor from the {x,y} which is the output of the fully connected layer from this post. My question is that how to do so that in the backward, I do not occur any problem??

Here is the sample code for the model.

def euclidean_dist(tensor1, tensor2):
     return distnace

class MyModel(): 

     def __init__():
          some convulotion layers...
          self.fc = nn.Linear(num_features, 3)

     def forward(in):

         (x, y, h) = fc(in)

         # Create the tensor with the axes value
         temp = torch.tensor([[a, b] for a in range(5) for b in range(5)])
         center = torch.tensor([[x,y]])

         # Now compute the distance between this and points
         distance = euclidean_distance(temp, centers)

         out = distance.view(5,5)
         out = out * h

         return out

Is it right to do this in this way?