diff --git a/docs/content/concepts/io-management/io-managers.mdx b/docs/content/concepts/io-management/io-managers.mdx
index 66331771a1062..0531fc7d21044 100644
--- a/docs/content/concepts/io-management/io-managers.mdx
+++ b/docs/content/concepts/io-management/io-managers.mdx
@@ -848,8 +848,9 @@ class DataframeTableIOManagerWithMetadata(ConfigurableIOManager):
         context.add_output_metadata({"num_rows": len(obj), "table_name": table_name})
 
     def load_input(self, context: InputContext):
-        table_name = context.upstream_output.name
-        return read_dataframe_from_table(name=table_name)
+        if context.upstream_output:
+            table_name = context.upstream_output.name
+            return read_dataframe_from_table(name=table_name)
 ```
 
 Any entries yielded this way will be attached to the `Handled Output` event for this output.
diff --git a/examples/docs_snippets/docs_snippets/concepts/io_management/custom_io_manager.py b/examples/docs_snippets/docs_snippets/concepts/io_management/custom_io_manager.py
index f6b3ce5a66d9e..93e3bb7e4e067 100755
--- a/examples/docs_snippets/docs_snippets/concepts/io_management/custom_io_manager.py
+++ b/examples/docs_snippets/docs_snippets/concepts/io_management/custom_io_manager.py
@@ -132,8 +132,9 @@ def handle_output(self, context: OutputContext, obj):
         context.add_output_metadata({"num_rows": len(obj), "table_name": table_name})
 
     def load_input(self, context: InputContext):
-        table_name = context.upstream_output.name
-        return read_dataframe_from_table(name=table_name)
+        if context.upstream_output:
+            table_name = context.upstream_output.name
+            return read_dataframe_from_table(name=table_name)
 
 
 # end_metadata_marker