# Buggy code
def main(args):
    # load and preprocess dataset
    data = load_data(args)
    features = torch.FloatTensor(data.features)
    labels = torch.LongTensor(data.labels)
    if hasattr(torch, "BoolTensor"):
        train_mask = torch.BoolTensor(data.train_mask)
        val_mask = torch.BoolTensor(data.val_mask)
        test_mask = torch.BoolTensor(data.test_mask)
    else:
        train_mask = torch.ByteTensor(data.train_mask)
        val_mask = torch.ByteTensor(data.val_mask)
        test_mask = torch.ByteTensor(data.test_mask)
    num_feats = features.shape[1]
    n_classes = data.num_labels
    n_edges = data.graph.number_of_edges()
    print(
        """----Data statistics------'
      #Edges %d
      #Classes %d 
      #Train samples %d
      #Val samples %d
      #Test samples %d"""
        % (
            n_edges,
            n_classes,
            train_mask.sum().item(),
            val_mask.sum().item(),
            test_mask.sum().item(),
        )
    )

    if args.gpu < 0:
        cuda = False
    else:
        cuda = True
        torch.cuda.set_device(args.gpu)
        features = features.cuda()
        labels = labels.cuda()
        train_mask = train_mask.cuda()
        val_mask = val_mask.cuda()
        test_mask = test_mask.cuda()

    g = data.graph
    # add self loop
    g.remove_edges_from(g.selfloop_edges())
    g = DGLGraph(g)
    g.add_edges(g.nodes(), g.nodes())
    n_edges = g.number_of_edges()
    # create model
    heads = ([args.num_heads] * args.num_layers) + [args.num_out_heads]
    model = GAT(
        g,
        args.num_layers,
        num_feats,
        args.num_hidden,
        n_classes,
        heads,
        F.elu,
        args.in_drop,
        args.attn_drop,
        args.negative_slope,
        args.residual,
    )
    print(model)
    stopper = EarlyStopping(patience=100)
    if cuda:
        model.cuda()
    loss_fcn = torch.nn.CrossEntropyLoss()

    # use optimizer
    optimizer = torch.optim.Adam(
        model.parameters(), lr=args.lr, weight_decay=args.weight_decay
    )

    # initialize graph
    dur = []
    for epoch in range(args.epochs):
        model.train()
        if epoch >= 3:
            t0 = time.time()
        # forward
        logits = model(features)
        loss = loss_fcn(logits[train_mask], labels[train_mask])

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if epoch >= 3:
            dur.append(time.time() - t0)

        train_acc = accuracy(logits[train_mask], labels[train_mask])

        if args.fastmode:
            val_acc = accuracy(logits[val_mask], labels[val_mask])
        else:
            val_acc = evaluate(model, features, labels, val_mask)
            if stopper.step(val_acc, model):
                break

        print(
            "Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | TrainAcc {:.4f} |"
            " ValAcc {:.4f} | ETputs(KTEPS) {:.2f}".format(
                epoch,
                np.mean(dur),
                loss.item(),
                train_acc,
                val_acc,
                n_edges / np.mean(dur) / 1000,
            )
        )

    print()
    model.load_state_dict(torch.load("es_checkpoint.pt"))
    acc = evaluate(model, features, labels, test_mask)
    print("Test Accuracy {:.4f}".format(acc))


# Traceback (AssertionError)
test_shared_mem_store.test_init ... /var/jenkins_home/workspace/DGL_PR-752@2/python/dgl/base.py:18: UserWarning: Initializer is not set. Use zero initializer instead. To suppress this warning, use `set_initializer` to explicitly specify which initializer to use.
warnings.warn(msg, warn_type)
/var/jenkins_home/workspace/DGL_PR-752@2/python/dgl/base.py:18: UserWarning: It may not be safe to access node data of all nodes.It's recommended to node data of a subset of nodes directly.
warnings.warn(msg, warn_type)
/var/jenkins_home/workspace/DGL_PR-752@2/python/dgl/base.py:18: UserWarning: It may not be safe to access edge data of all edges.It's recommended to edge data of a subset of edges directly.
warnings.warn(msg, warn_type)
[06:09:05] /var/jenkins_home/workspace/DGL_PR-752@2/src/runtime/shared_mem.cc:32: remove /test_graph1_node_feat for shared memory
[06:09:05] /var/jenkins_home/workspace/DGL_PR-752@2/src/runtime/shared_mem.cc:32: remove /test_graph1_node_test4 for shared memory
[06:09:05] /var/jenkins_home/workspace/DGL_PR-752@2/src/runtime/shared_mem.cc:32: remove /test_graph1_in for shared memory
[06:09:05] /var/jenkins_home/workspace/DGL_PR-752@2/src/runtime/shared_mem.cc:32: remove /test_graph1_edge_feat for shared memory
[06:09:05] /var/jenkins_home/workspace/DGL_PR-752@2/src/runtime/shared_mem.cc:32: remove /test_graph1_edge_test4 for shared memory
/var/jenkins_home/workspace/DGL_PR-752@2/python/dgl/base.py:18: UserWarning: Initializer is not set. Use zero initializer instead. To suppress this warning, use `set_initializer` to explicitly specify which initializer to use.
warnings.warn(msg, warn_type)
/var/jenkins_home/workspace/DGL_PR-752@2/python/dgl/base.py:18: UserWarning: It may not be safe to access node data of all nodes.It's recommended to node data of a subset of nodes directly.
warnings.warn(msg, warn_type)
/var/jenkins_home/workspace/DGL_PR-752@2/python/dgl/base.py:18: UserWarning: It may not be safe to access edge data of all edges.It's recommended to edge data of a subset of edges directly.
warnings.warn(msg, warn_type)

Arrays are not almost equal to 7 decimals

(mismatch 100.0%)
x: array([20., 20., 20., 20., 20., 20., 20., 20., 20., 20.], dtype=float32)
y: array(11)
Traceback (most recent call last):
File "/var/jenkins_home/workspace/DGL_PR-752@2/tests/distributed/test_shared_mem_store.py", line 65, in check_init_func
check_array_shared_memory(g, worker_id, [g.nodes[:].data['test4'], g.edges[:].data['test4']])
File "/var/jenkins_home/workspace/DGL_PR-752@2/tests/distributed/test_shared_mem_store.py", line 28, in check_array_shared_memory
assert_almost_equal(F.asnumpy(arr[0]), i + 10)
File "/usr/local/lib/python3.5/dist-packages/numpy/testing/nose_tools/utils.py", line 567, in assert_almost_equal
return assert_array_almost_equal(actual, desired, decimal, err_msg)
File "/usr/local/lib/python3.5/dist-packages/numpy/testing/nose_tools/utils.py", line 965, in assert_array_almost_equal
precision=decimal)
File "/usr/local/lib/python3.5/dist-packages/numpy/testing/nose_tools/utils.py", line 781, in assert_array_compare
raise AssertionError(msg)
AssertionError:
Arrays are not almost equal to 7 decimals

(mismatch 100.0%)
x: array([20., 20., 20., 20., 20., 20., 20., 20., 20., 20.], dtype=float32)
y: array(11)
FAIL

# Fixed code
def main(args):
    # load and preprocess dataset
    data = load_data(args)
    features = torch.FloatTensor(data.features)
    labels = torch.LongTensor(data.labels)
    if hasattr(torch, "BoolTensor"):
        train_mask = torch.BoolTensor(data.train_mask)
        val_mask = torch.BoolTensor(data.val_mask)
        test_mask = torch.BoolTensor(data.test_mask)
    else:
        train_mask = torch.ByteTensor(data.train_mask)
        val_mask = torch.ByteTensor(data.val_mask)
        test_mask = torch.ByteTensor(data.test_mask)
    num_feats = features.shape[1]
    n_classes = data.num_labels
    n_edges = data.graph.number_of_edges()
    print(
        """----Data statistics------'
      #Edges %d
      #Classes %d 
      #Train samples %d
      #Val samples %d
      #Test samples %d"""
        % (
            n_edges,
            n_classes,
            train_mask.sum().item(),
            val_mask.sum().item(),
            test_mask.sum().item(),
        )
    )

    if args.gpu < 0:
        cuda = False
    else:
        cuda = True
        torch.cuda.set_device(args.gpu)
        features = features.cuda()
        labels = labels.cuda()
        train_mask = train_mask.cuda()
        val_mask = val_mask.cuda()
        test_mask = test_mask.cuda()

    g = data.graph
    # add self loop
    g.remove_edges_from(nx.selfloop_edges(g))
    g = DGLGraph(g)
    g.add_edges(g.nodes(), g.nodes())
    n_edges = g.number_of_edges()
    # create model
    heads = ([args.num_heads] * args.num_layers) + [args.num_out_heads]
    model = GAT(
        g,
        args.num_layers,
        num_feats,
        args.num_hidden,
        n_classes,
        heads,
        F.elu,
        args.in_drop,
        args.attn_drop,
        args.negative_slope,
        args.residual,
    )
    print(model)
    stopper = EarlyStopping(patience=100)
    if cuda:
        model.cuda()
    loss_fcn = torch.nn.CrossEntropyLoss()

    # use optimizer
    optimizer = torch.optim.Adam(
        model.parameters(), lr=args.lr, weight_decay=args.weight_decay
    )

    # initialize graph
    dur = []
    for epoch in range(args.epochs):
        model.train()
        if epoch >= 3:
            t0 = time.time()
        # forward
        logits = model(features)
        loss = loss_fcn(logits[train_mask], labels[train_mask])

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if epoch >= 3:
            dur.append(time.time() - t0)

        train_acc = accuracy(logits[train_mask], labels[train_mask])

        if args.fastmode:
            val_acc = accuracy(logits[val_mask], labels[val_mask])
        else:
            val_acc = evaluate(model, features, labels, val_mask)
            if stopper.step(val_acc, model):
                break

        print(
            "Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | TrainAcc {:.4f} |"
            " ValAcc {:.4f} | ETputs(KTEPS) {:.2f}".format(
                epoch,
                np.mean(dur),
                loss.item(),
                train_acc,
                val_acc,
                n_edges / np.mean(dur) / 1000,
            )
        )

    print()
    model.load_state_dict(torch.load("es_checkpoint.pt"))
    acc = evaluate(model, features, labels, test_mask)
    print("Test Accuracy {:.4f}".format(acc))