How to use the wandb.watch function in wandb

def test_unwatch_multi(wandb_init_run):
    net1 = ConvNet()
    net2 = ConvNet()
    wandb.watch(net1, log_freq=1, log="all")
    wandb.watch(net2, log_freq=1, log="all")
    wandb.unwatch(net1)
    for i in range(3):
        output1 = net1(dummy_torch_tensor((64, 1, 28, 28)))
        output2 = net2(dummy_torch_tensor((64, 1, 28, 28)))
        grads = torch.ones(64, 10)
        output1.backward(grads)
        output2.backward(grads)
        assert(len(wandb_init_run.history.row) == 16)
        print(wandb_init_run.history.row)
        assert wandb_init_run.history.row.get('gradients/model_1/conv1.bias')
        assert wandb_init_run.history.row.get('gradients/conv1.bias') is None
        wandb.log({"a": 2})
    assert(len(wandb_init_run.history.rows) == 3)

def test_all_logging_freq(wandb_init_run):
    net = ConvNet()
    log_freq = 50
    wandb.watch(net, log="all", log_freq=log_freq)
    for i in range(110):
        output = net(dummy_torch_tensor((64, 1, 28, 28)))
        grads = torch.ones(64, 10)
        output.backward(grads)
        if (i + 1) % log_freq == 0:
            assert(len(wandb_init_run.history.row) == 16)
            assert(
                wandb_init_run.history.row['parameters/fc2.bias'].histogram[0] > 0)
            assert(
                wandb_init_run.history.row['gradients/fc2.bias'].histogram[0] > 0)
        else:
            assert(len(wandb_init_run.history.row) == 0)
        wandb.log({"a": 2})
    assert(len(wandb_init_run.history.rows) == 110)

def test_unwatch_multi(wandb_init_run):
    net1 = ConvNet()
    net2 = ConvNet()
    wandb.watch(net1, log_freq=1, log="all")
    wandb.watch(net2, log_freq=1, log="all")
    wandb.unwatch(net1)
    for i in range(3):
        output1 = net1(dummy_torch_tensor((64, 1, 28, 28)))
        output2 = net2(dummy_torch_tensor((64, 1, 28, 28)))
        grads = torch.ones(64, 10)
        output1.backward(grads)
        output2.backward(grads)
        assert(len(wandb_init_run.history.row) == 16)
        print(wandb_init_run.history.row)
        assert wandb_init_run.history.row.get('gradients/model_1/conv1.bias')
        assert wandb_init_run.history.row.get('gradients/conv1.bias') is None
        wandb.log({"a": 2})
    assert(len(wandb_init_run.history.rows) == 3)

def test_embedding(wandb_init_run):
    net = Embedding(d_embedding=300, d_word=300,
                    d_hidden=300, word_dim=100, dropout=0)
    wandb.watch(net, log="all", log_freq=1)
    for i in range(2):
        output = net(torch.ones((1, 4, 3, 224, 224)),
                     torch.ones((1, 4, 3, 20)))
        output.backward(torch.ones(1, 4, 300))
        wandb.log({"loss": 1})
    assert len(wandb_init_run.history.rows[0]) == 82

def train(self):
        """Train the agent."""
        # logger
        if self.args.log:
            wandb.init()
            wandb.config.update(self.hyper_params)
            wandb.watch([self.actor, self.critic], log="parameters")

        for i_episode in range(1, self.args.episode_num + 1):
            state = self.env.reset()
            done = False
            score = 0
            policy_loss_episode = list()
            value_loss_episode = list()
            self.episode_step = 0

            while not done:
                if self.args.render and i_episode >= self.args.render_after:
                    self.env.render()

                action = self.select_action(state)
                next_state, reward, done = self.step(action)
                policy_loss, value_loss = self.update_model()

# Initialize learning rate scheduler
    G_scheduler = None
    if args.lrsh is not None:
        G_scheduler = getattr(lr_scheduler, args.lrsh)(G_solver, gamma=args.lr_gamma_g)

    loss_fn = args.gen_g_loss[0]
    criterion_reconst = loss_function(loss_fn)

    # load saved models
    read_pickles_args(args, net, G_solver, net.vnet)

    # Initialize wandb
    if args.use_wandb:
        import wandb
        wandb.watch(net)

    try:
        iteration = G_solver.state_dict()['state'][G_solver.state_dict()['param_groups'][0]['params'][-1]]['step']
        epoch = int(iteration / int(dset_loaders_train.__len__() / args.batch_size))
        print('Continuing from iteration:{} epoch:{}'.format(iteration, epoch))
    except:
        epoch = 0
        iteration = 0
        print('Start training from scratch')

    num_view_samples = min(args.num_view_samples_per_batch, args.batch_size)

    if args.gen_g_loss[0] == 'weighted_dice':
        # todo: if implemented, remove the *box before the loss calculation
        raise NotImplementedError

def on_train_begin(self, **kwargs):
        "Call watch method to log model topology, gradients & weights"

        # Set self.best, method inherited from "TrackerCallback" by "SaveModelCallback"
        super().on_train_begin()

        # Ensure we don't call "watch" multiple times
        if not WandbCallback._watch_called:
            WandbCallback._watch_called = True

            # Logs model topology and optionally gradients and weights
            wandb.watch(self.learn.model, log=self.log)

if model_dir is not None:
        sfs_net_model.load_state_dict(torch.load(model_dir + 'sfs_net_model.pkl'))
    else:
        sfs_net_model.apply(weights_init)
        sfs_net_pretrained_dict = torch.load(pretrained_model_dict)
        sfs_net_state_dict = sfs_net_model.state_dict()
        load_model_from_pretrained(sfs_net_pretrained_dict, sfs_net_state_dict)
        sfs_net_model.load_state_dict(sfs_net_state_dict)
        sfs_net_model.fix_weights()

    os.system('mkdir -p {}'.format(args.log_dir))
    with open(args.log_dir+'/details.txt', 'w') as f:
        f.write(args.details)

    wandb.watch(sfs_net_model)
            
    # 1. Train on both Synthetic and Real (Celeba) dataset
    train(sfs_net_model, syn_data, celeba_data=celeba_data, read_first=read_first,\
            batch_size=batch_size, num_epochs=epochs, log_path=log_dir+'Mix_Training/', use_cuda=use_cuda, wandb=wandb, \
            lr=lr, wt_decay=wt_decay)

torch.nn.init.xavier_uniform_(m.weight)

dict_size = len(dictionary)
print("Dictionary size: ", dict_size)
print("Done reading babi!")

lstm = LSTM(args.hidden_dim_lstm, args.batch_size, dict_size, args.emb_dim, args.lstm_layers, device, dropout=args.dropout).to(device)
lstm.apply(init_weights)

input_dim_mlp = args.hidden_dim_lstm + args.hidden_dim_lstm + 40

rrn = RRN(args.hidden_dim_lstm, input_dim_mlp, args.hidden_dims_mlp, args.hidden_dim_rrn, args.message_dim_rrn, dict_size, args.f_dims, args.o_dims, device, args.batch_size, g_layers=1, edge_attribute_dim=args.hidden_dim_lstm, single_output=True, tanh=args.tanh_act, dropout=args.dropout).to(device)
rrn.apply(init_weights)

wandb.watch(lstm)
wandb.watch(rrn)

if args.load:
    load_models([(lstm, names_models[0]), (rrn, names_models[2])], result_folder, saving_path_rrn)

optimizer = torch.optim.Adam(chain(lstm.parameters(), rrn.parameters()), args.learning_rate, weight_decay=args.weight_decay)

criterion = torch.nn.CrossEntropyLoss(reduction='mean')

if args.epochs > 0:
    print("Start training")
    avg_train_losses, avg_train_accuracies, val_losses, val_accuracies = train(train_stories, validation_stories, args.epochs, lstm, rrn, criterion, optimizer, args.batch_size, args.no_save, device, result_folder)
    print("End training!")

if not args.test_on_test:
    test_stories = validation_stories

# Initialize models
    sfs_net_model      = SfsNetPipeline()
    if use_cuda:
        sfs_net_model = sfs_net_model.cuda()

    if model_dir is not None:
        sfs_net_model.load_state_dict(torch.load(model_dir + 'sfs_net_model.pkl'))
    else:
        print('Initializing weights')
        sfs_net_model.apply(weights_init)

    os.system('mkdir -p {}'.format(args.log_dir))
    with open(args.log_dir+'/details.txt', 'w') as f:
        f.write(args.details)

    wandb.watch(sfs_net_model)

    # 1. Train on Synthetic data
    train(sfs_net_model, syn_data, celeba_data = None, read_first=read_first, \
            batch_size=batch_size, num_epochs=epochs, log_path=log_dir+'Synthetic_Train/', use_cuda=use_cuda, wandb=wandb, \
            lr=lr, wt_decay=wt_decay)
    
    # 2. Generate Pseudo-Training information for CelebA dataset
    # Load CelebA dataset
    celeba_train_csv = celeba_data + '/train.csv'
    celeba_test_csv = celeba_data + '/test.csv'

    train_dataset, _ = get_celeba_dataset(read_from_csv=celeba_train_csv, read_first=read_first, validation_split=0)
    test_dataset, _ = get_celeba_dataset(read_from_csv=celeba_test_csv, read_first=read_first, validation_split=0)
    
    celeba_train_dl  = DataLoader(train_dataset, batch_size=1, shuffle=True)
    celeba_test_dl   = DataLoader(test_dataset, batch_size=1, shuffle=True)