Add PyTorch Memory Profiling Deck Renderer

flyteorg · Feb 2, 2025 · d7f5e51 · d7f5e51
1 parent 5607b0d
commit d7f5e51
Show file tree

Hide file tree

Showing 4 changed files with 1,161 additions and 0 deletions.
diff --git a/plugins/flytekit-deck-standard/flytekitplugins/deck/__init__.py b/plugins/flytekit-deck-standard/flytekitplugins/deck/__init__.py
@@ -14,6 +14,7 @@
    MarkdownRenderer
    SourceCodeRenderer
    TableRenderer
+   PyTorchProfilingRenderer
 """
 
 from .renderer import (
@@ -24,4 +25,16 @@
     MarkdownRenderer,
     SourceCodeRenderer,
     TableRenderer,
+    PyTorchProfilingRenderer
 )
+
+__all__ = [
+    "BoxRenderer",
+    "FrameProfilingRenderer",
+    "GanttChartRenderer",
+    "ImageRenderer",
+    "MarkdownRenderer",
+    "SourceCodeRenderer",
+    "TableRenderer",
+    "PyTorchProfilingRenderer"  
+]