def watch(models, criterion=None, log="gradients", log_freq=100):
    """
    Hooks into the torch model to collect gradients and the topology.  Should be extended
    to accept arbitrary ML models.

    :param (torch.Module) models: The model to hook, can be a tuple
    :param (torch.F) criterion: An optional loss value being optimized
    :param (str) log: One of "gradients", "parameters", "all", or None
    :param (int) log_freq: log gradients and parameters every N batches
    :return: (wandb.Graph) The graph object that will populate after the first backward pass
    """
    global watch_called
    if run is None:
        raise ValueError(
            "You must call `wandb.init` before calling watch")
    if watch_called:
        raise ValueError(
            "You can only call `wandb.watch` once per process. If you want to watch multiple models, pass them in as a tuple."
        )
    watch_called = True
    log_parameters = False
    log_gradients = True
    if log == "all":
        log_parameters = True
    elif log == "parameters":
        log_parameters = True
        log_gradients = False
    elif log is None:
        log_gradients = False

    if not isinstance(models, (tuple, list)):
        models = (models,)
    graphs = []
    prefix = ''
    for idx, model in enumerate(models):
        if idx > 0:
            prefix = "graph_%i" % idx

        run.history.torch.add_log_hooks_to_pytorch_module(
            model, log_parameters=log_parameters, log_gradients=log_gradients, prefix=prefix, log_freq=log_freq)

        graph = wandb_torch.TorchGraph.hook_torch(model, criterion, graph_idx=idx)
        graphs.append(graph)
        # NOTE: the graph is set in run.summary by hook_torch on the backward pass
    return graphs
Hooks into the torch model to collect gradients and the topology.  Should be extended
    to accept arbitrary ML models.

    :param (torch.Module) models: The model to hook, can be a tuple
    :param (torch.F) criterion: An optional loss value being optimized
    :param (str) log: One of "gradients", "parameters", "all", or None
    :param (int) log_freq: log gradients and parameters every N batches
    :return: (wandb.Graph) The graph object that will populate after the first backward pass