From 8de2449f84e2a7455a82736cf5c7f3a85fc88dcf Mon Sep 17 00:00:00 2001
From: Ke Wen <kw2501@meta.com>
Date: Mon, 3 Jun 2024 11:47:37 -0700
Subject: [PATCH] Update base for Update on "Add PP tracer + DP test"

Separate the addition of 2D test from original PR #362 for easier review and landing.

Also changed `.items()` to `named_children()` for more general submodule access. See similar changes in https://github.com/pytorch/torchtitan/pull/371.

[ghstack-poisoned]