diff --git a/src/workgroups2.el b/src/workgroups2.el
index 1ea55b59928f40bbaf616c2c70c0bc3fbcf6d756..72561897dc04d8181bcf83ac733aaca05f88150a 100644
--- a/src/workgroups2.el
+++ b/src/workgroups2.el
@@ -3307,7 +3307,9 @@ new workgroup during a switch.")
 
 (defun wg-flag-workgroup-modified (&optional workgroup)
   "Set WORKGROUP's and the current session's modified flags."
-  (when wg-flag-modified
+  (unless workgroup
+    (setq workgroup (wg-get-workgroup nil t)))
+  (when (and wg-flag-modified workgroup)
     (setf (wg-workgroup-modified workgroup) t)
     (wg-flag-session-modified)))