tensorrt_compatibility.patch

diff --git a/pid/_src/networks/pid_net.py b/pid/_src/networks/pid_net.py
index ab70f3b..8cffb22 100644
--- a/pid/_src/networks/pid_net.py
+++ b/pid/_src/networks/pid_net.py
@@ -461,7 +461,11 @@ class PidNet(PixDiT_T2I):
         # rank's local slice.
         if cp_group is not None:
             x_pixels = cat_outputs_cp_with_grad(x_pixels.contiguous(), seq_dim=2, cp_group=cp_group)
-        output = torch.nn.functional.fold(x_pixels, (H, W), kernel_size=self.patch_size, stride=self.patch_size)
+        
+        # Replace torch.nn.functional.fold with view and permute for TensorRT compatibility
+        x_pixels = x_pixels.view(B, C_out, self.patch_size, self.patch_size, Hs, Ws)
+        x_pixels = x_pixels.permute(0, 1, 4, 2, 5, 3).contiguous()
+        output = x_pixels.view(B, C_out, H, W)
 
         # Return (output, features) when feature extraction is enabled (without early exit)
         if feature_indices is not None and collected_features is not None:
diff --git a/pid/_src/networks/pixeldit_official.py b/pid/_src/networks/pixeldit_official.py
index a92670a..fb14757 100644
--- a/pid/_src/networks/pixeldit_official.py
+++ b/pid/_src/networks/pixeldit_official.py
@@ -63,11 +63,11 @@ def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
     assert embed_dim % 2 == 0
     if not torch.is_tensor(pos):
         pos = torch.as_tensor(pos)
-    omega = torch.arange(embed_dim // 2, dtype=torch.float64)
+    omega = torch.arange(embed_dim // 2, dtype=torch.float32)
     omega /= embed_dim / 2.0
     omega = 1.0 / 10000**omega  # (D/2,)
 
-    pos = pos.reshape(-1).to(torch.float64)  # (M,)
+    pos = pos.reshape(-1).to(torch.float32)  # (M,)
     out = torch.outer(pos, omega)  # (M, D/2), outer product
 
     emb_sin = torch.sin(out)  # (M, D/2)