Apply TOR901 only to files in test dir

[kit1980]

See pytorch/pytorch#125050 for details.

This will require passing file path as an input to the rules.
In general it doesn't make too much sense, but for PyTorch-internal rule like TOR901 where we know the path in the CI this will allow to remove long enumerations of files to ignore from .flake8

[sbrugman]

For reference, ruff has a option of opting out certain rules for a given pattern - which I think works quite well:

A list of mappings from file pattern to rule codes or prefixes to exclude, when considering any matching files. An initial '!' negates the file pattern.

https://docs.astral.sh/ruff/settings/#lint_per-file-ignores