From d18f2e6a07fedb5e6a0079f2fab60d8f516b2569 Mon Sep 17 00:00:00 2001 From: "Pyre Bot Jr." Date: Thu, 17 Jul 2025 12:55:04 -0700 Subject: [PATCH] Add missing Pyre mode headers] [batch:1/N] [A] [A] (#1016) Summary: X-link: https://github.com/pytorch/torchrec/pull/3204 X-link: https://github.com/pytorch/torchx/pull/1091 Pull Request resolved: https://github.com/pytorch/tnt/pull/1016 Differential Revision: D78138183 --- tests/utils/test_device_mesh.py | 2 ++ torchtnt/utils/device_mesh.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/tests/utils/test_device_mesh.py b/tests/utils/test_device_mesh.py index 1c486bdaac..2969e9d06d 100644 --- a/tests/utils/test_device_mesh.py +++ b/tests/utils/test_device_mesh.py @@ -4,6 +4,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +# pyre-strict + import unittest from torchtnt.utils.device_mesh import ( diff --git a/torchtnt/utils/device_mesh.py b/torchtnt/utils/device_mesh.py index 46e224ff23..8f343cffad 100644 --- a/torchtnt/utils/device_mesh.py +++ b/torchtnt/utils/device_mesh.py @@ -4,6 +4,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +# pyre-strict + from typing import Optional from torch.distributed.device_mesh import DeviceMesh, init_device_mesh