Repository: pytorch/ignite

Issue:
TensorboardLogger - tag/prefix for weights, grads and optimizer logging
### Problem
I have to log multiple models/optimizers with TensorboardLogger and I want them to be displayed in tensorboard in separate *subpanels* with distinct names/tags. 
At the moment, it is not possible to provide a custom tag or prefix for some handlers that log model parameters, because the tag is deduced from `model.named_parameters()`. 

As a consequence, parameters of multiple models are plottet to one graph.

### Solution
The affected handlers allow passing a custom prefix/tag to the function that is prepended to the deduced tag. 

E.g. for `GradsHistHandler` (snippet):
```python
logger.writer.add_histogram(tag="grads/{}".format(name),
                                        values=p.grad.detach().cpu().numpy(),
                                        global_step=global_step)
```

It can be changed to 
```python
logger.writer.add_histogram(tag=self.prefix + "grads/{}".format(name),
                                        values=p.grad.detach().cpu().numpy(),
                                        global_step=global_step)
```

Affected handlers in `ignite.contrib.handlers.tensorboard_logger`:
- ~~`OptimizerParamsHandler`~~
- `WeightsScalarHandler`
- `WeightsHistHandler`
- `GradsScalarHandler`
- `GradsHistHandler`


What do you think? 
If wanted, I can work on this. 




Discussion:
@Bibonaut I think we had something similar with `OptimizerParamsHandler` and we added `tag` in https://github.com/pytorch/ignite/pull/502

I agree that we may need to add `tag` to other handlers too.

So, yes, if you can send a PR on this, it would be very helpful !
@vfdev-5 Thank you for the hint with `OptimizerParamsHandler`.  I see that in #502 the `tag` property was also introduced in `BaseOptimizerParamsHandler`.  So I would do it the same way for gradients and weights handlers and add `tag` property to their base classes. Do you agree? 



@Bibonaut sounds good, please go ahead 👍 

Resolving patch:
diff --git a/ignite/contrib/handlers/base_logger.py b/ignite/contrib/handlers/base_logger.py
--- a/ignite/contrib/handlers/base_logger.py
+++ b/ignite/contrib/handlers/base_logger.py
@@ -130,7 +130,7 @@ class BaseWeightsScalarHandler(BaseHandler):
     Helper handler to log model's weights as scalars.
     """
 
-    def __init__(self, model, reduction=torch.norm):
+    def __init__(self, model, reduction=torch.norm, tag=None):
         if not isinstance(model, torch.nn.Module):
             raise TypeError("Argument model should be of type torch.nn.Module, "
                             "but given {}".format(type(model)))
@@ -149,6 +149,7 @@ def _is_0D_tensor(t):
 
         self.model = model
         self.reduction = reduction
+        self.tag = tag
 
 
 class BaseWeightsHistHandler(BaseHandler):
@@ -156,9 +157,10 @@ class BaseWeightsHistHandler(BaseHandler):
     Helper handler to log model's weights as histograms.
     """
 
-    def __init__(self, model):
+    def __init__(self, model, tag=None):
         if not isinstance(model, torch.nn.Module):
             raise TypeError("Argument model should be of type torch.nn.Module, "
                             "but given {}".format(type(model)))
 
         self.model = model
+        self.tag = tag
diff --git a/ignite/contrib/handlers/tensorboard_logger.py b/ignite/contrib/handlers/tensorboard_logger.py
--- a/ignite/contrib/handlers/tensorboard_logger.py
+++ b/ignite/contrib/handlers/tensorboard_logger.py
@@ -168,10 +168,11 @@ class WeightsScalarHandler(BaseWeightsScalarHandler):
     Args:
         model (torch.nn.Module): model to log weights
         reduction (callable): function to reduce parameters into scalar
+        tag (str, optional): common title for all produced plots. For example, 'generator'
 
     """
-    def __init__(self, model, reduction=torch.norm):
-        super(WeightsScalarHandler, self).__init__(model, reduction)
+    def __init__(self, model, reduction=torch.norm, tag=None):
+        super(WeightsScalarHandler, self).__init__(model, reduction, tag=tag)
 
     def __call__(self, engine, logger, event_name):
 
@@ -179,12 +180,13 @@ def __call__(self, engine, logger, event_name):
             raise RuntimeError("Handler 'WeightsScalarHandler' works only with TensorboardLogger")
 
         global_step = engine.state.get_event_attrib_value(event_name)
+        tag_prefix = "{}/".format(self.tag) if self.tag else ""
         for name, p in self.model.named_parameters():
             if p.grad is None:
                 continue
 
             name = name.replace('.', '/')
-            logger.writer.add_scalar("weights_{}/{}".format(self.reduction.__name__, name),
+            logger.writer.add_scalar("{}weights_{}/{}".format(tag_prefix, self.reduction.__name__, name),
                                      self.reduction(p.data),
                                      global_step)
 
@@ -208,23 +210,25 @@ class WeightsHistHandler(BaseWeightsHistHandler):
 
     Args:
         model (torch.nn.Module): model to log weights
+        tag (str, optional): common title for all produced plots. For example, 'generator'
 
     """
 
-    def __init__(self, model):
-        super(WeightsHistHandler, self).__init__(model)
+    def __init__(self, model, tag=None):
+        super(WeightsHistHandler, self).__init__(model, tag=tag)
 
     def __call__(self, engine, logger, event_name):
         if not isinstance(logger, TensorboardLogger):
             raise RuntimeError("Handler 'WeightsHistHandler' works only with TensorboardLogger")
 
         global_step = engine.state.get_event_attrib_value(event_name)
+        tag_prefix = "{}/".format(self.tag) if self.tag else ""
         for name, p in self.model.named_parameters():
             if p.grad is None:
                 continue
 
             name = name.replace('.', '/')
-            logger.writer.add_histogram(tag="weights/{}".format(name),
+            logger.writer.add_histogram(tag="{}weights/{}".format(tag_prefix, name),
                                         values=p.data.detach().cpu().numpy(),
                                         global_step=global_step)
 
@@ -251,22 +255,24 @@ class GradsScalarHandler(BaseWeightsScalarHandler):
     Args:
         model (torch.nn.Module): model to log weights
         reduction (callable): function to reduce parameters into scalar
+        tag (str, optional): common title for all produced plots. For example, 'generator'
 
     """
-    def __init__(self, model, reduction=torch.norm):
-        super(GradsScalarHandler, self).__init__(model, reduction)
+    def __init__(self, model, reduction=torch.norm, tag=None):
+        super(GradsScalarHandler, self).__init__(model, reduction, tag=tag)
 
     def __call__(self, engine, logger, event_name):
         if not isinstance(logger, TensorboardLogger):
             raise RuntimeError("Handler 'GradsScalarHandler' works only with TensorboardLogger")
 
         global_step = engine.state.get_event_attrib_value(event_name)
+        tag_prefix = "{}/".format(self.tag) if self.tag else ""
         for name, p in self.model.named_parameters():
             if p.grad is None:
                 continue
 
             name = name.replace('.', '/')
-            logger.writer.add_scalar("grads_{}/{}".format(self.reduction.__name__, name),
+            logger.writer.add_scalar("{}grads_{}/{}".format(tag_prefix, self.reduction.__name__, name),
                                      self.reduction(p.grad),
                                      global_step)
 
@@ -290,22 +296,24 @@ class GradsHistHandler(BaseWeightsHistHandler):
 
     Args:
         model (torch.nn.Module): model to log weights
+        tag (str, optional): common title for all produced plots. For example, 'generator'
 
     """
-    def __init__(self, model):
-        super(GradsHistHandler, self).__init__(model)
+    def __init__(self, model, tag=None):
+        super(GradsHistHandler, self).__init__(model, tag=tag)
 
     def __call__(self, engine, logger, event_name):
         if not isinstance(logger, TensorboardLogger):
             raise RuntimeError("Handler 'GradsHistHandler' works only with TensorboardLogger")
 
         global_step = engine.state.get_event_attrib_value(event_name)
+        tag_prefix = "{}/".format(self.tag) if self.tag else ""
         for name, p in self.model.named_parameters():
             if p.grad is None:
                 continue
 
             name = name.replace('.', '/')
-            logger.writer.add_histogram(tag="grads/{}".format(name),
+            logger.writer.add_histogram(tag="{}grads/{}".format(tag_prefix, name),
                                         values=p.grad.detach().cpu().numpy(),
                                         global_step=global_step)
 
diff --git a/ignite/contrib/handlers/visdom_logger.py b/ignite/contrib/handlers/visdom_logger.py
--- a/ignite/contrib/handlers/visdom_logger.py
+++ b/ignite/contrib/handlers/visdom_logger.py
@@ -234,11 +234,12 @@ class WeightsScalarHandler(BaseWeightsScalarHandler, _BaseVisDrawer):
     Args:
         model (torch.nn.Module): model to log weights
         reduction (callable): function to reduce parameters into scalar
+        tag (str, optional): common title for all produced plots. For example, 'generator'
         show_legend (bool, optional): flag to show legend in the window
     """
 
-    def __init__(self, model, reduction=torch.norm, show_legend=False):
-        super(WeightsScalarHandler, self).__init__(model, reduction)
+    def __init__(self, model, reduction=torch.norm, tag=None, show_legend=False):
+        super(WeightsScalarHandler, self).__init__(model, reduction, tag=tag)
         _BaseVisDrawer.__init__(self, show_legend=show_legend)
 
     def __call__(self, engine, logger, event_name):
@@ -247,9 +248,10 @@ def __call__(self, engine, logger, event_name):
             raise RuntimeError("Handler 'WeightsScalarHandler' works only with VisdomLogger")
 
         global_step = engine.state.get_event_attrib_value(event_name)
+        tag_prefix = "{}/".format(self.tag) if self.tag else ""
         for name, p in self.model.named_parameters():
             name = name.replace('.', '/')
-            k = "weights_{}/{}".format(self.reduction.__name__, name)
+            k = "{}weights_{}/{}".format(tag_prefix, self.reduction.__name__, name)
             v = float(self.reduction(p.data))
             self.add_scalar(logger, k, v, event_name, global_step)
 
@@ -278,12 +280,13 @@ class GradsScalarHandler(BaseWeightsScalarHandler, _BaseVisDrawer):
     Args:
         model (torch.nn.Module): model to log weights
         reduction (callable): function to reduce parameters into scalar
+        tag (str, optional): common title for all produced plots. For example, 'generator'
         show_legend (bool, optional): flag to show legend in the window
 
     """
 
-    def __init__(self, model, reduction=torch.norm, show_legend=False):
-        super(GradsScalarHandler, self).__init__(model, reduction)
+    def __init__(self, model, reduction=torch.norm, tag=None, show_legend=False):
+        super(GradsScalarHandler, self).__init__(model, reduction, tag)
         _BaseVisDrawer.__init__(self, show_legend=show_legend)
 
     def __call__(self, engine, logger, event_name):
@@ -291,9 +294,10 @@ def __call__(self, engine, logger, event_name):
             raise RuntimeError("Handler 'GradsScalarHandler' works only with VisdomLogger")
 
         global_step = engine.state.get_event_attrib_value(event_name)
+        tag_prefix = "{}/".format(self.tag) if self.tag else ""
         for name, p in self.model.named_parameters():
             name = name.replace('.', '/')
-            k = "grads_{}/{}".format(self.reduction.__name__, name)
+            k = "{}grads_{}/{}".format(tag_prefix, self.reduction.__name__, name)
             v = float(self.reduction(p.grad))
             self.add_scalar(logger, k, v, event_name, global_step)
 


Test patch:
diff --git a/tests/ignite/contrib/handlers/test_tensorboard_logger.py b/tests/ignite/contrib/handlers/test_tensorboard_logger.py
--- a/tests/ignite/contrib/handlers/test_tensorboard_logger.py
+++ b/tests/ignite/contrib/handlers/test_tensorboard_logger.py
@@ -230,23 +230,30 @@ def test_weights_scalar_handler(dummy_model_factory):
 
     model = dummy_model_factory(with_grads=True, with_frozen_layer=False)
 
-    wrapper = WeightsScalarHandler(model)
-    mock_logger = MagicMock(spec=TensorboardLogger)
-    mock_logger.writer = MagicMock()
+    # define test wrapper to test with and without optional tag
+    def _test(tag=None):
+        wrapper = WeightsScalarHandler(model, tag=tag)
+        mock_logger = MagicMock(spec=TensorboardLogger)
+        mock_logger.writer = MagicMock()
 
-    mock_engine = MagicMock()
-    mock_engine.state = State()
-    mock_engine.state.epoch = 5
+        mock_engine = MagicMock()
+        mock_engine.state = State()
+        mock_engine.state.epoch = 5
 
-    wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
+        wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
 
-    assert mock_logger.writer.add_scalar.call_count == 4
-    mock_logger.writer.add_scalar.assert_has_calls([
-        call("weights_norm/fc1/weight", 0.0, 5),
-        call("weights_norm/fc1/bias", 0.0, 5),
-        call("weights_norm/fc2/weight", 12.0, 5),
-        call("weights_norm/fc2/bias", math.sqrt(12.0), 5),
-    ], any_order=True)
+        tag_prefix = "{}/".format(tag) if tag else ""
+
+        assert mock_logger.writer.add_scalar.call_count == 4
+        mock_logger.writer.add_scalar.assert_has_calls([
+            call(tag_prefix + "weights_norm/fc1/weight", 0.0, 5),
+            call(tag_prefix + "weights_norm/fc1/bias", 0.0, 5),
+            call(tag_prefix + "weights_norm/fc2/weight", 12.0, 5),
+            call(tag_prefix + "weights_norm/fc2/bias", math.sqrt(12.0), 5),
+        ], any_order=True)
+
+    _test()
+    _test(tag="tag")
 
 
 def test_weights_scalar_handler_frozen_layers(dummy_model_factory):
@@ -294,23 +301,30 @@ def test_weights_hist_handler(dummy_model_factory):
 
     model = dummy_model_factory(with_grads=True, with_frozen_layer=False)
 
-    wrapper = WeightsHistHandler(model)
-    mock_logger = MagicMock(spec=TensorboardLogger)
-    mock_logger.writer = MagicMock()
+    # define test wrapper to test with and without optional tag
+    def _test(tag=None):
+        wrapper = WeightsHistHandler(model, tag=tag)
+        mock_logger = MagicMock(spec=TensorboardLogger)
+        mock_logger.writer = MagicMock()
 
-    mock_engine = MagicMock()
-    mock_engine.state = State()
-    mock_engine.state.epoch = 5
+        mock_engine = MagicMock()
+        mock_engine.state = State()
+        mock_engine.state.epoch = 5
 
-    wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
+        wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
 
-    assert mock_logger.writer.add_histogram.call_count == 4
-    mock_logger.writer.add_histogram.assert_has_calls([
-        call(tag="weights/fc1/weight", values=ANY, global_step=5),
-        call(tag="weights/fc1/bias", values=ANY, global_step=5),
-        call(tag="weights/fc2/weight", values=ANY, global_step=5),
-        call(tag="weights/fc2/bias", values=ANY, global_step=5),
-    ], any_order=True)
+        tag_prefix = "{}/".format(tag) if tag else ""
+
+        assert mock_logger.writer.add_histogram.call_count == 4
+        mock_logger.writer.add_histogram.assert_has_calls([
+            call(tag=tag_prefix + "weights/fc1/weight", values=ANY, global_step=5),
+            call(tag=tag_prefix + "weights/fc1/bias", values=ANY, global_step=5),
+            call(tag=tag_prefix + "weights/fc2/weight", values=ANY, global_step=5),
+            call(tag=tag_prefix + "weights/fc2/bias", values=ANY, global_step=5),
+        ], any_order=True)
+
+    _test()
+    _test(tag="tag")
 
 
 def test_weights_hist_handler_frozen_layers(dummy_model_factory):
@@ -359,25 +373,32 @@ def test_grads_scalar_handler_wrong_setup():
 def test_grads_scalar_handler(dummy_model_factory, norm_mock):
     model = dummy_model_factory(with_grads=True, with_frozen_layer=False)
 
-    wrapper = GradsScalarHandler(model, reduction=norm_mock)
-    mock_logger = MagicMock(spec=TensorboardLogger)
-    mock_logger.writer = MagicMock()
+    # define test wrapper to test with and without optional tag
+    def _test(tag=None):
+        wrapper = GradsScalarHandler(model, reduction=norm_mock, tag=tag)
+        mock_logger = MagicMock(spec=TensorboardLogger)
+        mock_logger.writer = MagicMock()
 
-    mock_engine = MagicMock()
-    mock_engine.state = State()
-    mock_engine.state.epoch = 5
-    norm_mock.reset_mock()
+        mock_engine = MagicMock()
+        mock_engine.state = State()
+        mock_engine.state.epoch = 5
+        norm_mock.reset_mock()
 
-    wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
+        wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
 
-    mock_logger.writer.add_scalar.assert_has_calls([
-        call("grads_norm/fc1/weight", ANY, 5),
-        call("grads_norm/fc1/bias", ANY, 5),
-        call("grads_norm/fc2/weight", ANY, 5),
-        call("grads_norm/fc2/bias", ANY, 5),
-    ], any_order=True)
-    assert mock_logger.writer.add_scalar.call_count == 4
-    assert norm_mock.call_count == 4
+        tag_prefix = "{}/".format(tag) if tag else ""
+
+        mock_logger.writer.add_scalar.assert_has_calls([
+            call(tag_prefix + "grads_norm/fc1/weight", ANY, 5),
+            call(tag_prefix + "grads_norm/fc1/bias", ANY, 5),
+            call(tag_prefix + "grads_norm/fc2/weight", ANY, 5),
+            call(tag_prefix + "grads_norm/fc2/bias", ANY, 5),
+        ], any_order=True)
+        assert mock_logger.writer.add_scalar.call_count == 4
+        assert norm_mock.call_count == 4
+
+    _test()
+    _test(tag="tag")
 
 
 def test_grads_scalar_handler_frozen_layers(dummy_model_factory, norm_mock):
@@ -424,23 +445,30 @@ def test_grads_hist_handler_wrong_setup():
 def test_grads_hist_handler(dummy_model_factory):
     model = dummy_model_factory(with_grads=True, with_frozen_layer=False)
 
-    wrapper = GradsHistHandler(model)
-    mock_logger = MagicMock(spec=TensorboardLogger)
-    mock_logger.writer = MagicMock()
+    # define test wrapper to test with and without optional tag
+    def _test(tag=None):
+        wrapper = GradsHistHandler(model, tag=tag)
+        mock_logger = MagicMock(spec=TensorboardLogger)
+        mock_logger.writer = MagicMock()
 
-    mock_engine = MagicMock()
-    mock_engine.state = State()
-    mock_engine.state.epoch = 5
+        mock_engine = MagicMock()
+        mock_engine.state = State()
+        mock_engine.state.epoch = 5
 
-    wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
+        wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
 
-    assert mock_logger.writer.add_histogram.call_count == 4
-    mock_logger.writer.add_histogram.assert_has_calls([
-        call(tag="grads/fc1/weight", values=ANY, global_step=5),
-        call(tag="grads/fc1/bias", values=ANY, global_step=5),
-        call(tag="grads/fc2/weight", values=ANY, global_step=5),
-        call(tag="grads/fc2/bias", values=ANY, global_step=5),
-    ], any_order=True)
+        tag_prefix = "{}/".format(tag) if tag else ""
+
+        assert mock_logger.writer.add_histogram.call_count == 4
+        mock_logger.writer.add_histogram.assert_has_calls([
+            call(tag=tag_prefix + "grads/fc1/weight", values=ANY, global_step=5),
+            call(tag=tag_prefix + "grads/fc1/bias", values=ANY, global_step=5),
+            call(tag=tag_prefix + "grads/fc2/weight", values=ANY, global_step=5),
+            call(tag=tag_prefix + "grads/fc2/bias", values=ANY, global_step=5),
+        ], any_order=True)
+
+    _test()
+    _test(tag="tag")
 
 
 def test_grads_hist_frozen_layers(dummy_model_factory):
diff --git a/tests/ignite/contrib/handlers/test_visdom_logger.py b/tests/ignite/contrib/handlers/test_visdom_logger.py
--- a/tests/ignite/contrib/handlers/test_visdom_logger.py
+++ b/tests/ignite/contrib/handlers/test_visdom_logger.py
@@ -390,34 +390,44 @@ def __init__(self):
 
     model = DummyModel()
 
-    wrapper = WeightsScalarHandler(model)
-    mock_logger = MagicMock(spec=VisdomLogger)
-    mock_logger.vis = MagicMock()
-    mock_logger.executor = _DummyExecutor()
+    # define test wrapper to test with and without optional tag
+    def _test(tag=None):
+        wrapper = WeightsScalarHandler(model, tag=tag)
+        mock_logger = MagicMock(spec=VisdomLogger)
+        mock_logger.vis = MagicMock()
+        mock_logger.executor = _DummyExecutor()
 
-    mock_engine = MagicMock()
-    mock_engine.state = State()
-    mock_engine.state.epoch = 5
-
-    wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
+        mock_engine = MagicMock()
+        mock_engine.state = State()
+        mock_engine.state.epoch = 5
 
-    assert mock_logger.vis.line.call_count == 4
-    mock_logger.vis.line.assert_has_calls([
-        call(X=[5, ], Y=[0.0, ], env=mock_logger.vis.env,
-             win=None, update=None,
-             opts=wrapper.windows["weights_norm/fc1/weight"]['opts'], name="weights_norm/fc1/weight"),
-        call(X=[5, ], Y=[0.0, ], env=mock_logger.vis.env,
-             win=None, update=None,
-             opts=wrapper.windows["weights_norm/fc1/bias"]['opts'], name="weights_norm/fc1/bias"),
-
-        call(X=[5, ], Y=[12.0, ], env=mock_logger.vis.env,
-             win=None, update=None,
-             opts=wrapper.windows["weights_norm/fc2/weight"]['opts'], name="weights_norm/fc2/weight"),
-        call(X=[5, ], Y=ANY, env=mock_logger.vis.env,
-             win=None, update=None,
-             opts=wrapper.windows["weights_norm/fc2/bias"]['opts'], name="weights_norm/fc2/bias"),
+        wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
 
-    ], any_order=True)
+        tag_prefix = "{}/".format(tag) if tag else ""
+
+        assert mock_logger.vis.line.call_count == 4
+        mock_logger.vis.line.assert_has_calls([
+            call(X=[5, ], Y=[0.0, ], env=mock_logger.vis.env,
+                 win=None, update=None,
+                 opts=wrapper.windows[tag_prefix + "weights_norm/fc1/weight"]['opts'],
+                 name=tag_prefix + "weights_norm/fc1/weight"),
+            call(X=[5, ], Y=[0.0, ], env=mock_logger.vis.env,
+                 win=None, update=None,
+                 opts=wrapper.windows[tag_prefix + "weights_norm/fc1/bias"]['opts'],
+                 name=tag_prefix + "weights_norm/fc1/bias"),
+
+            call(X=[5, ], Y=[12.0, ], env=mock_logger.vis.env,
+                 win=None, update=None,
+                 opts=wrapper.windows[tag_prefix + "weights_norm/fc2/weight"]['opts'],
+                 name=tag_prefix + "weights_norm/fc2/weight"),
+            call(X=[5, ], Y=ANY, env=mock_logger.vis.env,
+                 win=None, update=None,
+                 opts=wrapper.windows[tag_prefix + "weights_norm/fc2/bias"]['opts'],
+                 name=tag_prefix + "weights_norm/fc2/bias"),
+        ], any_order=True)
+
+    _test()
+    _test(tag="tag")
 
 
 def test_weights_scalar_handler_custom_reduction():
@@ -502,35 +512,44 @@ def __init__(self):
     def norm(x):
         return 0.0
 
-    wrapper = GradsScalarHandler(model, reduction=norm)
-    mock_logger = MagicMock(spec=VisdomLogger)
-    mock_logger.vis = MagicMock()
-    mock_logger.executor = _DummyExecutor()
+    # define test wrapper to test with and without optional tag
+    def _test(tag=None):
+        wrapper = GradsScalarHandler(model, reduction=norm, tag=tag)
+        mock_logger = MagicMock(spec=VisdomLogger)
+        mock_logger.vis = MagicMock()
+        mock_logger.executor = _DummyExecutor()
 
-    mock_engine = MagicMock()
-    mock_engine.state = State()
-    mock_engine.state.epoch = 5
-
-    wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
+        mock_engine = MagicMock()
+        mock_engine.state = State()
+        mock_engine.state.epoch = 5
 
-    assert mock_logger.vis.line.call_count == 4
-
-    mock_logger.vis.line.assert_has_calls([
-        call(X=[5, ], Y=ANY, env=mock_logger.vis.env,
-             win=None, update=None,
-             opts=wrapper.windows["grads_norm/fc1/weight"]['opts'], name="grads_norm/fc1/weight"),
-        call(X=[5, ], Y=ANY, env=mock_logger.vis.env,
-             win=None, update=None,
-             opts=wrapper.windows["grads_norm/fc1/bias"]['opts'], name="grads_norm/fc1/bias"),
-
-        call(X=[5, ], Y=ANY, env=mock_logger.vis.env,
-             win=None, update=None,
-             opts=wrapper.windows["grads_norm/fc2/weight"]['opts'], name="grads_norm/fc2/weight"),
-        call(X=[5, ], Y=ANY, env=mock_logger.vis.env,
-             win=None, update=None,
-             opts=wrapper.windows["grads_norm/fc2/bias"]['opts'], name="grads_norm/fc2/bias"),
+        wrapper(mock_engine, mock_logger, Events.EPOCH_STARTED)
 
-    ], any_order=True)
+        tag_prefix = "{}/".format(tag) if tag else ""
+
+        assert mock_logger.vis.line.call_count == 4
+        mock_logger.vis.line.assert_has_calls([
+            call(X=[5, ], Y=ANY, env=mock_logger.vis.env,
+                 win=None, update=None,
+                 opts=wrapper.windows[tag_prefix + "grads_norm/fc1/weight"]['opts'],
+                 name=tag_prefix + "grads_norm/fc1/weight"),
+            call(X=[5, ], Y=ANY, env=mock_logger.vis.env,
+                 win=None, update=None,
+                 opts=wrapper.windows[tag_prefix + "grads_norm/fc1/bias"]['opts'],
+                 name=tag_prefix + "grads_norm/fc1/bias"),
+
+            call(X=[5, ], Y=ANY, env=mock_logger.vis.env,
+                 win=None, update=None,
+                 opts=wrapper.windows[tag_prefix + "grads_norm/fc2/weight"]['opts'],
+                 name=tag_prefix + "grads_norm/fc2/weight"),
+            call(X=[5, ], Y=ANY, env=mock_logger.vis.env,
+                 win=None, update=None,
+                 opts=wrapper.windows[tag_prefix + "grads_norm/fc2/bias"]['opts'],
+                 name=tag_prefix + "grads_norm/fc2/bias"),
+        ], any_order=True)
+
+    _test()
+    _test(tag="tag")
 
 
 def test_integration_no_server():