Register modules set to None in config for test_save_load_optional_components

Author: clarencechen

tests/test_pipelines_common.py

@@ -394,9 +394,10 @@ def test_save_load_optional_components(self):
         pipe.to(torch_device)
         pipe.set_progress_bar_config(disable=None)
 
-        # set all optional components to None
+        # set all optional components to None and update pipeline config accordingly
         for optional_component in pipe._optional_components:
             setattr(pipe, optional_component, None)
+        pipe.register_modules(**{optional_component: None for optional_component in pipe._optional_components})
 
         inputs = self.get_dummy_inputs(torch_device)
         output = pipe(**inputs)[0]