@driazati  I added a helper function and wrote this code:
@torch.jit.script
def center_slice_helper(layer, diff_y, diff_x, h_end, w_end):
   return layer[:, :, diff_y:h_end, diff_x:w_end]

def center_crop(self, layer, target_size):
    #_, _, layer_height, layer_width = layer.size()

    diff_y = (layer.shape[2] - target_size.shape[2]) // 2
    diff_x = (layer.shape[3] - target_size.shape[3]) // 2

    h_end = diff_y + target_size.shape[2]
    w_end = diff_x + target_size.shape[3]

    return center_slice_helper(layer, diff_y, diff_x, h_end, w_end)


def forward(self, x, bridge):

    up = self.up(x)
    crop1 = self.center_crop(bridge, up)
    out = torch.cat([up, crop1], 1)
    out = self.conv_block(out)

    return out

At this time, it gave this exception:
torch.jit.TracingCheckError: Tracing failed sanity checks!
Encountered an exception while running the Python function with test inputs.
Exception:
	center_slice_helper() expected value of type Tensor for argument 'diff_y' in position 1, but instead got value of type int.
	Value: 0
	Declaration: center_slice_helper(Tensor layer, Tensor diff_y, Tensor diff_x, Tensor h_end, Tensor w_end) -> Tensor

I debugged the code and observed that x parameter of forward function has this type:

x_parameter.PNG932×390 8.66 KB

How to handle this?