Reduce the memory leak caused by torch.jit.trace

[ydshieh]

What does this PR do?

Reduce the memory leak caused by torch.jit.trace.

Without this fix, each call to _create_and_check_torchscript increases RAM usage by ~20MB.
(Even with this call, there are still memory leak by ~0.04MB)

Remark

Since our torchscript tests are slow tests, they are run only in scheduled CI (where the test jobs are organized by models), this memory leak issue is not important (memory is released after each job process is done/exit).

However, for PR like #16679, I need to make sure the modified tests will pass. When I run in a GCP VM, I got OOM issue.
So the change in this PR still have its value, I think.

More information

The method is copied from torch

https:/pytorch/pytorch/blob/bcf6974c207ac0339bfb8bdfdb0b0ec348f7a22f/torch/testing/_internal/jit_utils.py#L59

This is called in torch/test, including

https:/pytorch/pytorch/blob/bcf6974c207ac0339bfb8bdfdb0b0ec348f7a22f/test/test_jit.py#L12962

    # clear the class registry as we will be defining foo multiple times
    jit_utils.clear_class_registry()

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[LysandreJik]

Oh that's a very good catch, and may well have been one of the causes of the out-of-memory errors we were getting when running all tests sequentially.

[sgugger]

Thanks for fixing this!