diff --git a/src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts b/src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts
index ad3d3319d5..85457f167a 100644
--- a/src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts
+++ b/src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts
@@ -257,6 +257,7 @@ export abstract class NotebookInput extends EditorInput implements INotebookInpu
 		this._notebookEditorOpenedTimestamp = Date.now();
 		if (this._textInput) {
 			this.hookDirtyListener(this._textInput.onDidChangeDirty, () => this._onDidChangeDirty.fire());
+			this._register(this._textInput);
 		}
 	}