From c69939aa85460e8135f40bce908a6cddaa73065f Mon Sep 17 00:00:00 2001
From: Edward Yang <ezyang@meta.com>
Date: Thu, 21 Nov 2024 11:08:51 -0800
Subject: [PATCH] Prepare for "Fix type-safety of `torch.nn.Module` instances":
 wave 2

Summary: See D52890934

Reviewed By: malfet, r-barnes

Differential Revision: D66245100

fbshipit-source-id: 019058106ac7eaacf29c1c55912922ea55894d23
---
 .../DensePose/densepose/modeling/losses/cycle_pix2shape.py    | 2 ++
 .../DensePose/densepose/modeling/losses/cycle_shape2shape.py  | 2 ++
 projects/DensePose/densepose/modeling/losses/embed.py         | 4 ++++
 projects/DensePose/densepose/modeling/losses/soft_embed.py    | 4 ++++
 projects/DensePose/densepose/modeling/roi_heads/roi_head.py   | 2 ++
 5 files changed, 14 insertions(+)

diff --git a/projects/DensePose/densepose/modeling/losses/cycle_pix2shape.py b/projects/DensePose/densepose/modeling/losses/cycle_pix2shape.py
index f4e0a94a68..0fc4298ffa 100644
--- a/projects/DensePose/densepose/modeling/losses/cycle_pix2shape.py
+++ b/projects/DensePose/densepose/modeling/losses/cycle_pix2shape.py
@@ -149,6 +149,8 @@ def forward(
         return torch.stack(losses, dim=0).mean()
 
     def fake_value(self, densepose_predictor_outputs: Any, embedder: nn.Module):
+        # pyre-fixme[29]: `Union[(self: Tensor) -> Any, Module, Tensor]` is not a
+        #  function.
         losses = [embedder(mesh_name).sum() * 0 for mesh_name in embedder.mesh_names]
         losses.append(densepose_predictor_outputs.embedding.sum() * 0)
         return torch.mean(torch.stack(losses))
diff --git a/projects/DensePose/densepose/modeling/losses/cycle_shape2shape.py b/projects/DensePose/densepose/modeling/losses/cycle_shape2shape.py
index c6e2ca89a3..0eb7a8ee3a 100644
--- a/projects/DensePose/densepose/modeling/losses/cycle_shape2shape.py
+++ b/projects/DensePose/densepose/modeling/losses/cycle_shape2shape.py
@@ -61,6 +61,8 @@ def forward(self, embedder: nn.Module):
 
     def fake_value(self, embedder: nn.Module):
         losses = []
+        # pyre-fixme[29]: `Union[(self: Tensor) -> Any, Module, Tensor]` is not a
+        #  function.
         for mesh_name in embedder.mesh_names:
             losses.append(embedder(mesh_name).sum() * 0)
         return torch.mean(torch.stack(losses))
diff --git a/projects/DensePose/densepose/modeling/losses/embed.py b/projects/DensePose/densepose/modeling/losses/embed.py
index 2fc8a16a47..20790073b8 100644
--- a/projects/DensePose/densepose/modeling/losses/embed.py
+++ b/projects/DensePose/densepose/modeling/losses/embed.py
@@ -104,6 +104,8 @@ def __call__(
             ) / (-self.embdist_gauss_sigma)
             losses[mesh_name] = F.cross_entropy(scores, vertex_indices_i, ignore_index=-1)
 
+        # pyre-fixme[29]: `Union[(self: Tensor) -> Any, Module, Tensor]` is not a
+        #  function.
         for mesh_name in embedder.mesh_names:
             if mesh_name not in losses:
                 losses[mesh_name] = self.fake_value(
@@ -113,6 +115,8 @@ def __call__(
 
     def fake_values(self, densepose_predictor_outputs: Any, embedder: nn.Module):
         losses = {}
+        # pyre-fixme[29]: `Union[(self: Tensor) -> Any, Module, Tensor]` is not a
+        #  function.
         for mesh_name in embedder.mesh_names:
             losses[mesh_name] = self.fake_value(densepose_predictor_outputs, embedder, mesh_name)
         return losses
diff --git a/projects/DensePose/densepose/modeling/losses/soft_embed.py b/projects/DensePose/densepose/modeling/losses/soft_embed.py
index f746d67a75..1ace3dac30 100644
--- a/projects/DensePose/densepose/modeling/losses/soft_embed.py
+++ b/projects/DensePose/densepose/modeling/losses/soft_embed.py
@@ -118,6 +118,8 @@ def __call__(
             )
             losses[mesh_name] = (-geodist_softmax_values * embdist_logsoftmax_values).sum(1).mean()
 
+        # pyre-fixme[29]: `Union[(self: Tensor) -> Any, Module, Tensor]` is not a
+        #  function.
         for mesh_name in embedder.mesh_names:
             if mesh_name not in losses:
                 losses[mesh_name] = self.fake_value(
@@ -127,6 +129,8 @@ def __call__(
 
     def fake_values(self, densepose_predictor_outputs: Any, embedder: nn.Module):
         losses = {}
+        # pyre-fixme[29]: `Union[(self: Tensor) -> Any, Module, Tensor]` is not a
+        #  function.
         for mesh_name in embedder.mesh_names:
             losses[mesh_name] = self.fake_value(densepose_predictor_outputs, embedder, mesh_name)
         return losses
diff --git a/projects/DensePose/densepose/modeling/roi_heads/roi_head.py b/projects/DensePose/densepose/modeling/roi_heads/roi_head.py
index 66997270d8..d274f2f412 100644
--- a/projects/DensePose/densepose/modeling/roi_heads/roi_head.py
+++ b/projects/DensePose/densepose/modeling/roi_heads/roi_head.py
@@ -158,6 +158,7 @@ def _forward_densepose(self, features: Dict[str, torch.Tensor], instances: List[
                 proposal_boxes = [x.proposal_boxes for x in proposals]
 
                 if self.use_decoder:
+                    # pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
                     features_list = [self.decoder(features_list)]
 
                 features_dp = self.densepose_pooler(features_list, proposal_boxes)
@@ -171,6 +172,7 @@ def _forward_densepose(self, features: Dict[str, torch.Tensor], instances: List[
             pred_boxes = [x.pred_boxes for x in instances]
 
             if self.use_decoder:
+                # pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
                 features_list = [self.decoder(features_list)]
 
             features_dp = self.densepose_pooler(features_list, pred_boxes)