diff --git a/src/components/notebook/create-assignment/creation-switch.tsx b/src/components/notebook/create-assignment/creation-switch.tsx
index 88feaee..9118e49 100644
--- a/src/components/notebook/create-assignment/creation-switch.tsx
+++ b/src/components/notebook/create-assignment/creation-switch.tsx
@@ -86,7 +86,13 @@ export class CreationModeSwitch extends React.Component<IModeSwitchProps> {
         } else {
           currentLayout.widgets.map(w => {
             if (w instanceof CreationWidget || w instanceof ErrorWidget) {
-              currentLayout.removeWidget(w);
+              try {
+                currentLayout.removeWidget(w);
+              } catch(error: any) {
+                console.log("Could not remove widget: " + w)
+                console.log("Error: " + error)
+              }
+              
             }
           });
         }
diff --git a/src/components/notebook/manual-grading/grading-switch.tsx b/src/components/notebook/manual-grading/grading-switch.tsx
index c22ee1a..867ca55 100644
--- a/src/components/notebook/manual-grading/grading-switch.tsx
+++ b/src/components/notebook/manual-grading/grading-switch.tsx
@@ -160,7 +160,12 @@ export class GradingModeSwitch extends React.Component<IModeSwitchProps> {
         } else {
           currentLayout.widgets.map(w => {
             if (w instanceof DataWidget || w instanceof GradeWidget) {
-              currentLayout.removeWidget(w);
+              try {
+                currentLayout.removeWidget(w);
+              } catch(error: any) {
+                console.log("Could not remove widget:" + w)
+                console.log("Error: " + error)
+              }
             }
           });
         }