Expose layer

torch-ext/deformable_detr/__init__.py

@@ -2,6 +2,7 @@ from typing import List
 import torch
 
 from ._ops import ops
+from . import layers
 
 
 def ms_deform_attn_backward(
@@ -42,4 +43,4 @@ def ms_deform_attn_forward(
     )
 
 
-__all__ = ["ms_deform_attn_forward", "ms_deform_attn_backward"]
+__all__ = ["layers", "ms_deform_attn_forward", "ms_deform_attn_backward"]

torch-ext/deformable_detr/layers.py

@@ -79,3 +79,6 @@ class MultiScaleDeformableAttention(nn.Module):
             attention_weights,
             im2col_step,
         )
+
+
+__all__ = ["MultiScaleDeformableAttention"]