Apply TOR901 only to files in test dir

[kit1980]

See https://github.com/pytorch/pytorch/issues/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