[train] bump test_torch_trainer timeout

[matthewdeng]

Description

Bump from small to medium due to timeouts happening specifically in py3.12 tests.

[gemini-code-assist]

Code Review

This pull request addresses test timeouts for test_torch_trainer on Python 3.12 by increasing the test size from small to medium. This is a valid approach. However, I've suggested a more precise solution using Bazel's select functionality to apply the larger timeout only when building with Python 3.12. This would avoid increasing the timeout for other Python versions, contributing to better CI efficiency. See my comment for an example implementation.

[gemini-code-assist]

A more targeted fix for the Python 3.12 timeouts would be to use Bazel's select statement. This would apply the medium size only for the Python 3.12 build, while other versions would continue to use the small size, which can help keep CI runs fast.

Here's an example of how you could implement this:

First, define a config_setting (if one doesn't already exist for your toolchain):

# In a relevant BUILD file (e.g., the root BUILD file)
config_setting(
    name = "is_py312",
    flag_values = {
        "@bazel_tools//tools/python:python_version": "PY3.12",
    },
)

Then, use select in this py_test definition:

py_test(
    name = "test_torch_trainer",
    size = select({
        "//path/to:is_py312": "medium",
        "//conditions:default": "small",
    }),
    ...
)