## 3.6 Decentralized Stochastic Gradient Descent

In previous sections, we have discussed AWC-DGD and ATC-DGD algorithms. In this section, we will discuss their stochastic variants.

### 3.6.1 The stochastic optimization problem

Consider the optimization problem:

$$\min_{x \in \mathbb{R}^d} \quad f(x) = \frac{1}{n}\sum_{i=1}^n [f_i(x) = \mathbb{E}_{\xi_i \sim D_i}\{F(x,\xi_i)\}]$$

where $f_i(x)$ is local to node $i$, and random variable $\xi_i$ denotes the local data that follows distribution $D_i$. Each node $i$ can locally evaluate stochastic gradient $\nabla F(x;\xi_i)$ (not the real gradient $\nabla f_i(x)$); it must communicate to access information from other nodes. Since the real gradient $\nabla f_i(x)$ cannot be accessed, we cannot use ATC-DGD or AWC-DGD introduced in the above sections to solve it. The above stochastic optimization problem is very common in large-scale signal processing, machine learning, and deep learning.

### 3.6.2 Adaptation-then-combination decentralized stochastic gradient descent (ATC-DSGD)

Given the network topology and the weight matrix, the ATC-DSGD will update as follows:

\begin{align}
x_i^{(k+1)} =  \sum_{j=1}^n w_{ij} \Big(x_j^{(k)} - \alpha \nabla F(x_j^{(k)};\xi_j^{(k)}) \Big), \quad  i=1,\cdots,n.
\end{align}

Comparing the above recursion with ATC-DGD, the only difference is to replace the real gradient $\nabla f_i(x)$ with the stochastic gradient $\nabla F(x_i^{(k)};\xi_i^{(k)})$, where $\xi_i^{(k)}$ is a realization of the random variable $\xi_i$ at iteration $k$.

### 3.6.3 Adaptation-with-combination decentralized stochastic gradient descent (AWC-DSGD)

Given the network topology and the weight matrix, the AWC-DSGD will update as follows:

\begin{align}
x_i^{(k+1)} =  \sum_{j=1}^n w_{ij} x_j^{(k)} - \alpha \nabla F(x_i^{(k)};\xi_i^{(k)}), \quad  i=1,\cdots,n.
\end{align}

Comparing the above recursion with AWC-DGD, the only difference is to replace the real gradient $\nabla f_i(x)$ with the stochastic gradient $\nabla F(x_i^{(k)};\xi_i^{(k)})$. Since AWC-DSGD shares the same algorithm structure with ATC-DSGD, it can be conducted in a non-blocking manner in which the computation and communication steps can be carried out in an independent and parallel manner.

### 3.6.4 Convergence property

#### 3.6.4.1  Convergence of ATC-DSGD

Now we discuss the convergence property of ATC-DSGD and AWC-DSGD. The following assumptions are quite standard in literatures.

**A.1** Each $f_i(x)$ is $L$-smooth and $\mu$-strongly convex.

**A.2** The random sample  $\xi_i^{(k)}$ is independent of each other for any $k$ and $i$. We also assume each stochastic gradient is unbiased and has bounded variance, i.e., $\mathbb{E}[{\nabla}F(x;\xi_{i})] = \nabla f_i(x)$ and $\mathbb{E}\|{\nabla}F(x;\xi_{i}) - \nabla f_i(x) \|^2 \le  \sigma^2$.

**A.3** The network topology is strongly connected, and the weight matrix $W$ is doubly-stochastic.

With the above assumptions, it is established in \[YAYS-2020\] that

\begin{align}\label{dsgd-convergence}
\frac{1}{n}\sum_{i=1}^n \mathbb{E}\|x_i^{(k)} - x^\star\|^2 = O\Big( \underbrace{(1-\alpha \mu)^k}_{\rm convg.\ rate} + \underbrace{\frac{\alpha \sigma^2}{n} + \frac{\rho^2 \alpha^2 \sigma^2}{1-\rho}}_{\rm sto.\ bias} + \underbrace{\frac{\rho^2 \alpha^2 b^2}{(1-\rho)^2}}_{\rm inconsis.\ bias} \Big). \hspace{1cm} \mbox{(ATC-DSGD-Convergence)}
\end{align}

where $x^\star$ is the glboal solution to the optimization problem, $\rho = \max\{|\lambda_2(W), \lambda_n(W)|\}$ and $b^2 = \frac{1}{n}\sum_{i=1}^n \|\nabla f_i(x^\star)\|^2$ denotes the data heterogeneity between nodes. Quantity $1-\rho$ measures the connectivity of the network topology. It is observed that AWC-DGD cannot converge exactly to the solution $x^\star$, but to a neighborhood around it. The limiting error is on the order of $O(\frac{\alpha^2 b^2}{(1-\rho)^2})$. When step-size $\alpha$ is small, or the data heterogeneity $b^2$ is small, or the network is well-connected, i.e., $\rho \to 0$, the limiting error can be negligible.

It is observed from (ATC-DSGD-Convergence) that ATC-DSGD suffers from two sorces of bias: the stochastic gradient bias, and the inconsistency bias. When the real gradient $\nabla f_i(x)$ can be accessed, it holds that $\sigma^2 = 0$ and (ATC-DSGD-Convergence) will reduce to (ATC-DGD-Convergence) in Sec. 3.2. Observating (ATC-DSGD-Convergence), we know that:

- When step-size $\alpha$ is sufficiently small, we have $\lim_{k\to \infty} \frac{1}{n}\sum_{i=1}^n \mathbb{E}\|x_i^{(k)} - x^\star\|^2 = O(\frac{\alpha \sigma^2}{n})$, which is determined by the gradient noise $\sigma^2$.

- When the network is badly-connected such that $\rho \to 0$, we have $\lim_{k\to \infty} \frac{1}{n}\sum_{i=1}^n \mathbb{E}\|x_i^{(k)} - x^\star\|^2 = O(\frac{\alpha^2 b^2}{(1-\rho)^2})$, which is determined by the data heterogeneity $b^2$.

#### 3.6.4.2  Convergence of AWC-DSGD

With the above assumptions, it holds that

\begin{align}
\frac{1}{n}\sum_{i=1}^n \mathbb{E}\|x_i^{(k)} - x^\star\|^2 = O\Big( \underbrace{(1-\alpha \mu)^k}_{\rm convg.\ rate} + \underbrace{\frac{\alpha \sigma^2}{n} + \frac{\alpha^2 \sigma^2}{1-\rho}}_{\rm sto.\ bias} + \underbrace{\frac{\alpha^2 b^2}{(1-\rho)^2}}_{\rm inconsis.\ bias} \Big). \hspace{1cm} \mbox{(AWC-DSGD-Convergence)}
\end{align}

Comparing (AWC-DSGD-Convergence) and (ATC-DSGD-Convergence), we find ATC-DSGD can achieve a better accuracy when the same step-size and network topology are used for both algorithms becasue $\rho \in (0,1)$. Moreover, AWC-DSGD is less stable than ATC-DSGD, see \[Sayed-2014, Secs.10.6 and 11.5\].

### 3.6.5 Application: Logistic regression for MNIST

In this subsection, we will employ ATC-DSGD to coduct a binary classification task using MNIST dataset. In particular, we will utilize PyTorch to compute the gradient automatically rather than compute the gradient by ourselves. The following experiment has two purposes:

1. To illutrate how DSGD works for distributed optimization

2. To illustrate how to implement DSGD with BlueFog, and how it can be easily correlates with PyTorch AutoGrad.

#### 3.6.5.1 Distributed logistic regression

Consider the following logistic regression problem:

$$\min_{w\in \mathbb{R}^d} \quad \frac{1}{n}\sum_{i=1}^n \mathbb{E}_{y_i,h_i}[\ln(1+\exp(-y_i x_i^T w))] + \frac{\rho}{2}\|w\|^2$$

where $\{x_i, y_i\}$ are local data held in node $i$, $\rho$ is a small positive constant.

#### 3.6.5.1 Test iparallel setup

```
import ipyparallel as ipp
import numpy as np
import torch
import networkx as nx  # nx will be used for network topology creating and plotting
import matplotlib.pyplot as plt

%matplotlib inline
```

```
rc = ipp.Client(profile="bluefog")
num_nodes = len(rc.ids)
dview = rc[:]  # A DirectView of all engines
dview.block = True
rc.ids
```

    [0, 1, 2, 3]

```
%%px
import numpy as np
import bluefog.torch as bf
import torch
from bluefog.common import topology_util
import networkx as nx
from sklearn.datasets import make_classification
from sklearn.preprocessing import normalize

bf.init()
print(f"Hello, I am {bf.rank()} among {bf.size()} processes")
```

    [stdout:0] Hello, I am 1 among 4 processes
    [stdout:1] Hello, I am 3 among 4 processes
    [stdout:2] Hello, I am 2 among 4 processes
    [stdout:3] Hello, I am 0 among 4 processes

#### 3.6.5.2 Prepare data

```
# Load the MNIST dataset.

from torchvision import datasets

dataset = datasets.MNIST("../data", train=True, download=True)
```

Since logistic regression is a binary problem we will take a subset of the data that has labels 0 or 1. In addition, because MNIST images are 28 $\times$ 28 images, we will reshape each image to be a vector of size 784. Each image in the MNIST data set is grayscale so an element of a MNIST image ranges from 0 to 255.

```
idx = (dataset.targets == 0) | (dataset.targets == 1)
X = dataset.data[idx].view(dataset.data[idx].shape[0], -1)
y = dataset.targets[idx][:, None]
```

We will confirm the shape of our new data and label matrix.

```
print("Feature X has the shape:", X.shape)
print("Label y has the shape:", y.shape)
print("Lable y should only contains labels 0 and 1:", y[:10].t())
```

    Feature X has the shape: torch.Size([12665, 784])
    Label y has the shape: torch.Size([12665, 1])
    Lable y should only contains labels 0 and 1: tensor([[0, 1, 1, 1, 1, 0, 1, 1, 0, 0]])

We will normalize each image such that each element will range from 0 to 1 instead of 0 to 255.

```
X = X / 255.0
```

The data $(X,y)$ is ready. We can plot two example images.

```
plt.imshow(X[0, :].view(28, 28))
```

```
plt.imshow(X[1, :].view(28, 28))
```

Next we partion data and move each partition to the node.

```
# Move the data to each worker.
M, dim = X.shape
num_split = int(M // num_nodes)
for i in range(num_nodes):
    X_worker = X[i * num_split : (i + 1) * num_split]
    y_worker = y[i * num_split : (i + 1) * num_split]
    y_worker = torch.squeeze(y_worker)
    y_worker = y_worker.type(torch.FloatTensor)
    #     X_worker = X_worker.to(torch.double)
    #     y_worker = y_worker.to(torch.double)
    dview.push({"X_worker": X_worker, "y_worker": y_worker}, targets=i)
```

    /Users/ybc/anaconda3/envs/bf_test/lib/python3.7/site-packages/torch/storage.py:34: FutureWarning: pickle support for Storage will be removed in 1.5. Use `torch.save` instead
      warnings.warn("pickle support for Storage will be removed in 1.5. Use `torch.save` instead", FutureWarning)

Confirm that the data has been moved successfully.

```
%%px
print(X_worker.shape)
print(y_worker.shape)
```

    [stdout:0]
    torch.Size([3166, 784])
    torch.Size([3166])
    [stdout:1]
    torch.Size([3166, 784])
    torch.Size([3166])
    [stdout:2]
    torch.Size([3166, 784])
    torch.Size([3166])
    [stdout:3]
    torch.Size([3166, 784])
    torch.Size([3166])

Each local dataset $(X_i, y_i)$ at worker $i$ is ready.

#### 3.6.5.3 Train the binary classifier with DSGD

```
%%px

# Implement functions for logistic regression.

def sigmoid(x):
    """Return the output of the sigmoid function.
    Args:
     - x (torch tensor): Input that we want to apply the sigmoid function on.
    Return:
    - Return a torch tensor of the output

    """
    return 1.0 / (1.0 + torch.exp(-x))

def obj(w, X, y, reg):
    """Output of the objective function of logistic regression.
    Args:
     - w (torch tensor): Parameter.
     - X (torch tensor): Data.
     - y (torch tensor): Labels.
     - reg (float): Regularization parameter.
    Return:
     - Returns a float of the value of the objective function.
    """
    M, _ = X.shape

    return -1 / M * torch.sum(
        y[:, None] * X @ w - torch.log(1 + torch.exp(X @ w))
    ) + reg / 2 * torch.norm(w, p=2)

def grad(w, X, y, reg):
    """Output of the gradient for the logistic regression.
    Args:
     - w (torch tensor): Parameter.
     - X_batch (torch tensor): Batch of data gradient is calculated on.
     - y_batch (torch tensor): Labels.
     - reg (float): Regularization parameter.
    Return:
     - Returns a torch tensor that is the gradient for the data and parameters provided.
    """
    M, _ = X.shape
    return -1 / M * (X.T @ y[:, None] - (X.T @ (1 - sigmoid(-X @ w)))) + reg * w
```

```
%%px
from ipyparallel.datapub import publish_data

# Create a distributed gradient descent function that is equivalent to centralized gradient descent.
def distributed_grad_descent(X, y, reg=1e-4, maxite=1000, step_size=1):
    M, dim = X.shape
    w_opt = 0.001 * torch.randn((dim, 1))
    for i in range(maxite):
        # Calculate local gradient.
        grad_iter = grad(w_opt, X, y, reg)
        # Use local gradient to calculate global gradient.
        gradient = bf.allreduce(grad_iter)

        # Distributed gradient descent.
        w_opt = w_opt - step_size * gradient
        if i % 20 == 0:
            publish_data(dict(i=i, gradient_norm=torch.norm(gradient, p=2)))
    publish_data(
        dict(i=maxite - 1, gradient_norm=torch.norm(gradient, p=2), w_opt=w_opt)
    )
    print(f"[DG] Rank {bf.rank()}: Global Gradient Norm: {torch.norm(gradient, p=2)}")
    return w_opt
```

```
%%px

# Adapt then combine using blocking communication.

def ATC(X, y, w_opt, w=None, step_size=1.0, reg=1e-4):
    """
    Adapt then combine algorithm using blocking communication. Performs one iteration.
    Inputs:
    - X (torch tensor) := The data to calculate the loss on of size (N, dim).
    - y (torch tensor) := The corresponding labels of the data of size (N, 1).
    - w_opt (torch tensor) := The optimal parameters.
    - w (torch tensor) := The parameter to update over.
    - step_size (float) := The step size parameter.
    - reg (float) := The regularization parameter.
    - batch_size (int) := Size of a batch for the iteration.
    Output:
    - w (torch tensor) := Update iteration.
    - Relative error.
    """
    M, dim = X.shape

    if w is None:
        w = 0.001 * torch.randn((dim, 1))

    adapt = w - step_size * grad(w, X, y, reg)
    combine = bf.neighbor_allreduce(adapt)
    w = combine

    dist = torch.norm(w - w_opt, p=2) / torch.norm(w_opt, p=2)
    # averaged_relative_error = bf.allreduce(dist**2, average=True)

    return w, dist  # torch.sqrt(averaged_relative_error)
```

```
%%px

# Adapt with combine using non-blocking communication.

def AWC(X, y, w_opt, w=None, step_size=1.0, reg=1e-4):
    """
    Adapt with combine algorithm using non-blocking communication. Performs one iteration.
    Inputs:
    - X (torch tensor) := The data to calculate the loss on of size (N, dim).
    - y (torch tensor) := The corresponding labels of the data of size (N, 1).
    - w_opt (torch tensor) := The optimal parameters.
    - w (torch tensor) := The parameter to update over.
    - step_size (float) := The step size parameter.
    - reg (float) := The regularization parameter.
    - batch_size (int) := Size of a batch for the iteration.
    Output:
    - w (torch tensor) := Update iteration.
    - Relative error.
    """
    M, dim = X.shape

    if w is None:
        w = 0.001 * torch.randn((dim, 1))

    combine_handle = bf.neighbor_allreduce_nonblocking(w)
    adapt = step_size * grad(w, X, y, reg)
    combine = bf.wait(combine_handle)
    w = combine - adapt

    dist = torch.norm(w - w_opt, p=2) / torch.norm(w_opt, p=2)
    # averaged_relative_error = bf.allreduce(dist**2, average=True)

    return w, dist  # torch.sqrt(averaged_relative_error)
```

```
from tqdm import tqdm

maxite = 5000
dview["maxite"] = maxite
ar = dview.apply_async(
    lambda: distributed_grad_descent(X_worker, y_worker, maxite=maxite)
)
i = 0

# This may take a while finish it.
with tqdm(total=maxite - 1, desc="Epoch     #{}".format(i + 1)) as t:
    while not ar.ready():
        data = ar.data[0]
        t.set_postfix({"gradient_norm": data.get("gradient_norm", 1)})
        t.update(data.get("i", 0) - t.n)
w_opt = ar.data[0]["w_opt"]
dview["w_opt"] = w_opt
```

    Iter     #1: 100%|██████████| 499/499 [00:06<00:00, 79.22it/s, gradient_norm=tensor(0.0015)]

```
plt.imshow(w_opt.view(28, 28))
```

```
%%px
from ipyparallel.datapub import publish_data

# Create a distributed gradient descent function that is equivalent to centralized gradient descent.
def decentralized_grad_descent(
    one_step_func, X, y, w_opt, maxite, step_size=1.0, reg=1e-4
):
    print(maxite)
    M, dim = X.shape
    w = 0.001 * torch.randn((dim, 1))
    for i in range(maxite):
        w, relative_error = one_step_func(
            X, y, w=w, w_opt=w_opt, step_size=step_size, reg=reg
        )
        if i % 20 == 0:
            publish_data(dict(i=i, relative_error=relative_error, w=w))
    publish_data(dict(i=maxite - 1, relative_error=relative_error, w=w))
```

```
%%px
decentralized_grad_descent(ATC, X_worker, y_worker, w_opt, maxite=100)
```

    [stdout:0] 100
    [stdout:1] 100
    [stdout:2] 100
    [stdout:3] 100

```
maxite = 500
dview["maxite"] = maxite
ar = dview.apply_async(
    lambda: decentralized_grad_descent(ATC, X_worker, y_worker, w_opt, maxite=maxite)
)
i = 0

# This may take a while finish it.
with tqdm(total=maxite - 1, desc="Epoch     #{}".format(i + 1)) as t:
    while not ar.ready():
        data = ar.data[0]
        t.set_postfix({"relative_error": data.get("relative_error", 1)})
        t.update(data.get("i", 0) - t.n)
        if data.get("i", 0) != 0:
            w = data.get("w", None)
            plt.imshow(w.view(28, 28))
            plt.title(f"i={data.get('i',0)}")
            plt.show()
w = ar.data[0]["w"]
dview["w"] = w
```

```
%%px

# loss_step
def ATC_DGD(X, y, rho=1e-2, alpha=1e-3, epoch=50):

    n, d = X.shape  # n:local data size, d:feature dimension
    w = torch.zeros(d, 1, dtype=torch.double, requires_grad=True)

    # Set topology as exponential-two topology.
    G = topology_util.ExponentialTwoGraph(bf.size())
    bf.set_topology(G)

    # main recursion
    for _ in range(epoch):
        shuffle_order = torch.randperm(n)
        for order in shuffle_order[:100]:

            # sample one data
            i = order.item()
            feature = X[i].reshape(1, -1)
            label = y[i]

            # compute grad automatically
            loss_ = torch.mean(
                torch.log(1 + torch.exp(-label * feature.mm(w)))
            ) + 0.5 * rho * torch.norm(w, p=2)
            loss_.backward()

            # ATC-DGD
            with torch.no_grad():
                w.data = bf.neighbor_allreduce(w.data - alpha * w.grad.data)
                w.grad.data.zero_()

    return w
```

```
%%px

w = ATC_DGD(X_worker, y_worker, rho=1e-2, alpha=1e-3, epoch=5)
```

```
1 + 1
```

    2

Find the optimal result using the distributed gradient descent function developed before.

```
print(w_opt_mnist)
```

    [stdout:0]
    [DG] Rank 0: Global Gradient Norm: 0.00021530537924263626
    tensor([[-4.1454e-04],
            [ 1.5268e-03],
            [ 2.2789e-03],
            [-5.2180e-04],
            [-2.4208e-04],
            [-8.2633e-04],
            [-1.1322e-04],
            [-2.4679e-04],
            [-4.1106e-04],
            [ 3.2596e-04],
            [-4.6936e-04],
            [ 6.6684e-05],
            [-5.1959e-04],
            [ 1.7390e-04],
            [-4.0965e-04],
            [-1.1527e-03],
            [-1.3719e-03],
            [ 6.6315e-04],
            [-4.6811e-05],
            [-4.3492e-04],
            [ 7.3475e-05],
            [ 1.5381e-03],
            [-3.8010e-04],
            [-6.2565e-04],
            [ 6.3816e-04],
            [-1.5623e-03],
            [-3.8932e-04],
            [-1.4504e-03],
            [-9.6471e-05],
            [ 7.3598e-04],
            [ 2.3206e-03],
            [-9.5610e-04],
            [ 4.7032e-04],
            [-8.5588e-04],
            [-1.2766e-03],
            [-1.1936e-03],
            [ 8.5900e-04],
            [-2.3481e-06],
            [ 1.3068e-03],
            [ 1.3673e-03],
            [-7.2730e-05],
            [-8.5659e-04],
            [ 1.1147e-03],
            [-2.0424e-03],
            [-1.2847e-02],
            [-6.2633e-03],
            [ 9.7928e-04],
            [ 1.8101e-03],
            [ 6.1258e-04],
            [-1.5042e-03],
            [ 1.3360e-04],
            [-1.7237e-04],
            [ 2.0086e-04],
            [ 1.0413e-04],
            [ 3.5428e-04],
            [-7.8566e-04],
            [ 8.3567e-05],
            [ 1.0233e-03],
            [ 1.3289e-04],
            [-8.7346e-04],
            [-1.8640e-03],
            [ 7.6127e-04],
            [ 6.6773e-04],
            [ 2.2841e-03],
            [ 9.0441e-04],
            [-2.6830e-04],
            [-1.3008e-03],
            [ 2.4222e-04],
            [ 7.8573e-03],
            [ 1.0450e-01],
            [ 3.2059e-02],
            [ 1.4406e-02],
            [ 6.0007e-02],
            [ 1.3585e-01],
            [ 3.4749e-02],
            [ 3.0322e-03],
            [-6.2479e-04],
            [ 7.1583e-04],
            [-1.6865e-04],
            [ 3.3282e-04],
            [ 3.5418e-04],
            [ 7.0774e-04],
            [ 3.1967e-04],
            [ 1.0994e-03],
            [ 1.1855e-03],
            [ 3.2760e-04],
            [-1.0656e-03],
            [ 1.5669e-03],
            [ 1.7766e-03],
            [-1.2481e-03],
            [ 8.4661e-04],
            [-3.8667e-04],
            [ 7.6094e-03],
            [ 1.0458e-02],
            [ 7.1351e-04],
            [ 2.4327e-03],
            [ 5.6904e-02],
            [ 1.1743e-01],
            [ 1.1777e-01],
            [ 1.2267e-01],
            [ 2.3519e-01],
            [ 2.6924e-01],
            [ 1.2881e-01],
            [ 1.8973e-02],
            [ 9.7233e-03],
            [ 3.6758e-03],
            [-1.6747e-03],
            [-7.6741e-04],
            [-5.3438e-04],
            [ 8.2690e-04],
            [-4.8871e-04],
            [-1.6149e-04],
            [-3.2138e-04],
            [ 1.5634e-03],
            [-3.4550e-05],
            [-9.1970e-04],
            [ 2.2506e-04],
            [ 9.9152e-04],
            [-6.7695e-06],
            [-5.2260e-03],
            [ 8.4027e-02],
            [ 1.3167e-01],
            [ 4.4042e-02],
            [ 4.9748e-02],
            [ 1.1187e-01],
            [ 1.2508e-01],
            [ 5.5261e-02],
            [ 1.2392e-01],
            [ 7.3954e-02],
            [ 2.0312e-01],
            [-9.3808e-02],
            [ 3.5205e-02],
            [ 1.5864e-01],
            [ 1.0079e-01],
            [ 7.6359e-02],
            [ 3.7612e-02],
            [ 5.9002e-03],
            [-4.1301e-04],
            [-5.5714e-04],
            [ 1.5475e-03],
            [ 1.0044e-03],
            [-9.7317e-04],
            [-9.0777e-04],
            [ 1.5642e-03],
            [ 1.4684e-03],
            [ 4.2056e-03],
            [-1.6220e-03],
            [-7.9440e-03],
            [ 1.2386e-01],
            [ 1.4461e-01],
            [ 3.2836e-02],
            [ 4.6390e-02],
            [ 1.7389e-01],
            [ 2.8755e-01],
            [-4.5606e-02],
            [-1.4891e-01],
            [-5.3263e-04],
            [ 6.7824e-02],
            [-1.4261e-01],
            [-5.6060e-02],
            [ 2.1152e-01],
            [ 1.8650e-01],
            [ 1.4992e-01],
            [ 1.3038e-01],
            [ 5.4771e-02],
            [ 1.0132e-03],
            [ 4.8791e-04],
            [ 7.9906e-04],
            [-6.8106e-04],
            [ 1.0756e-03],
            [ 1.5110e-03],
            [ 2.1278e-03],
            [ 6.9284e-04],
            [-1.1194e-03],
            [ 1.7695e-02],
            [ 7.3869e-02],
            [ 8.0635e-02],
            [ 2.5154e-02],
            [-3.9285e-02],
            [-6.2712e-02],
            [ 7.5195e-02],
            [ 1.0874e-01],
            [-4.9581e-02],
            [-1.1885e-01],
            [-1.4442e-01],
            [-6.6793e-02],
            [-5.4998e-02],
            [-1.1059e-01],
            [ 2.1092e-01],
            [ 2.1775e-01],
            [ 7.9351e-03],
            [ 3.1830e-02],
            [ 4.3387e-02],
            [ 2.1414e-04],
            [-1.6228e-03],
            [-6.2277e-04],
            [-1.3110e-03],
            [ 3.5542e-06],
            [-1.8492e-04],
            [ 6.1647e-04],
            [ 1.3138e-01],
            [ 1.4127e-01],
            [ 6.3568e-02],
            [ 1.4701e-01],
            [ 1.9911e-02],
            [-1.5686e-01],
            [-2.4967e-01],
            [-2.8631e-01],
            [-4.7523e-02],
            [ 4.5806e-02],
            [-6.5777e-03],
            [-3.1996e-02],
            [ 1.2006e-02],
            [-1.7620e-01],
            [-1.3366e-01],
            [-1.4290e-01],
            [ 5.4982e-02],
            [-5.2510e-02],
            [-1.0864e-01],
            [-5.7900e-02],
            [-6.2358e-02],
            [-9.7890e-03],
            [-1.5302e-04],
            [ 4.7228e-04],
            [-9.9241e-04],
            [-3.1786e-04],
            [ 7.4502e-04],
            [ 7.0709e-05],
            [ 1.2193e-01],
            [ 1.6224e-01],
            [ 5.9834e-02],
            [ 1.1051e-01],
            [ 1.0461e-02],
            [-1.6331e-01],
            [-2.0378e-01],
            [-1.4479e-01],
            [ 1.4055e-01],
            [ 2.5251e-03],
            [ 2.4745e-01],
            [ 1.5639e-01],
            [-2.0631e-01],
            [-2.5072e-01],
            [ 1.3546e-03],
            [-1.6667e-02],
            [-1.9821e-01],
            [-2.7018e-01],
            [-1.5122e-01],
            [-6.2037e-02],
            [-3.3747e-02],
            [-1.8915e-03],
            [-1.8294e-03],
            [ 6.9204e-04],
            [-1.5357e-04],
            [ 1.6708e-03],
            [ 1.0733e-03],
            [-1.0558e-03],
            [ 1.4822e-01],
            [ 1.3818e-01],
            [-4.1531e-02],
            [ 4.7091e-02],
            [ 1.9003e-02],
            [-3.9035e-02],
            [-8.6894e-02],
            [-8.1036e-03],
            [ 1.2747e-01],
            [ 2.7214e-01],
            [ 2.9709e-01],
            [ 1.1567e-01],
            [-1.5539e-02],
            [-9.5740e-02],
            [ 1.6482e-01],
            [ 1.2793e-01],
            [-2.8419e-01],
            [-3.9789e-01],
            [-2.4394e-01],
            [-1.3566e-01],
            [-2.6044e-02],
            [-7.2220e-03],
            [-4.1549e-04],
            [-1.2107e-03],
            [ 1.0983e-04],
            [ 4.6884e-05],
            [-2.2104e-03],
            [-5.3214e-04],
            [ 1.4655e-01],
            [ 4.1023e-02],
            [-1.2377e-01],
            [-1.2125e-01],
            [-1.1221e-01],
            [-1.5344e-02],
            [-3.6775e-02],
            [ 4.2039e-02],
            [ 1.1091e-01],
            [ 1.1478e-01],
            [ 1.1573e-01],
            [ 2.6739e-01],
            [ 2.1596e-01],
            [-1.0234e-01],
            [ 1.3185e-02],
            [-2.1221e-02],
            [-3.0426e-01],
            [-3.4172e-01],
            [-2.8324e-01],
            [-1.5865e-01],
            [-4.5446e-02],
            [-5.5689e-03],
            [ 1.2696e-03],
            [ 4.8764e-04],
            [-2.4698e-04],
            [-6.5265e-04],
            [ 5.6271e-04],
            [ 3.2708e-04],
            [ 1.3862e-01],
            [-3.7834e-02],
            [-1.6532e-01],
            [-1.9285e-01],
            [-2.2579e-01],
            [-1.6731e-01],
            [-1.2137e-01],
            [-9.5400e-02],
            [-5.1763e-03],
            [ 4.8719e-02],
            [ 1.2580e-01],
            [ 5.0455e-01],
            [ 3.3722e-01],
            [-2.5439e-02],
            [-1.7518e-01],
            [-2.5705e-01],
            [-1.1951e-01],
            [-2.5577e-01],
            [-2.5881e-01],
            [-9.0948e-02],
            [-7.8044e-02],
            [-4.3442e-03],
            [ 1.2903e-04],
            [-6.0843e-04],
            [ 1.4594e-03],
            [ 5.7002e-04],
            [-1.3896e-03],
            [-6.9921e-04],
            [ 7.6237e-02],
            [-1.2488e-01],
            [-1.7576e-01],
            [-1.6374e-01],
            [-2.7732e-01],
            [-2.9823e-01],
            [-1.2015e-01],
            [-1.3605e-01],
            [-2.3567e-01],
            [ 1.6882e-01],
            [ 3.7708e-01],
            [ 6.8420e-01],
            [ 3.1567e-01],
            [-7.1413e-03],
            [-3.2854e-01],
            [-3.2361e-01],
            [-7.2232e-02],
            [-2.9126e-01],
            [-2.6734e-01],
            [-1.7813e-01],
            [-9.4485e-02],
            [-1.0271e-02],
            [-1.5277e-03],
            [-8.6602e-04],
            [-8.1627e-04],
            [-1.3140e-04],
            [ 3.2578e-06],
            [ 2.5387e-04],
            [ 3.4196e-03],
            [-1.6792e-01],
            [-1.7129e-01],
            [-2.7093e-01],
            [-3.4530e-01],
            [-2.2278e-01],
            [-1.7824e-01],
            [-3.4777e-01],
            [-1.6606e-01],
            [ 4.6749e-01],
            [ 7.1757e-01],
            [ 6.3815e-01],
            [ 2.7185e-01],
            [-7.5921e-02],
            [-6.6452e-01],
            [-2.4247e-01],
            [ 2.1679e-03],
            [-4.0710e-01],
            [-4.2250e-01],
            [-2.6234e-01],
            [-1.0597e-01],
            [-1.3102e-02],
            [ 8.1200e-04],
            [ 1.6037e-03],
            [-1.5470e-03],
            [ 5.9793e-04],
            [-4.3401e-04],
            [ 1.3630e-04],
            [-3.8539e-02],
            [-1.7989e-01],
            [-2.5330e-01],
            [-3.6467e-01],
            [-2.5319e-01],
            [-6.5559e-02],
            [-4.1015e-01],
            [-3.1432e-01],
            [-2.0814e-02],
            [ 5.3430e-01],
            [ 8.9337e-01],
            [ 5.5051e-01],
            [ 2.9156e-01],
            [-2.5857e-01],
            [-8.2349e-01],
            [-1.4299e-01],
            [-1.4367e-01],
            [-4.0789e-01],
            [-3.4698e-01],
            [-2.6304e-01],
            [-1.1774e-01],
            [-1.4676e-02],
            [-4.1625e-04],
            [-2.5845e-04],
            [ 4.2557e-06],
            [ 1.5651e-03],
            [ 1.0527e-03],
            [ 5.7057e-04],
            [-1.0840e-01],
            [-2.2952e-01],
            [-2.7911e-01],
            [-2.9994e-01],
            [-2.1410e-01],
            [-2.0011e-01],
            [-4.9769e-01],
            [-2.4322e-01],
            [ 2.0222e-01],
            [ 7.5507e-01],
            [ 9.3063e-01],
            [ 3.3976e-01],
            [ 3.0932e-02],
            [-5.7397e-01],
            [-4.9022e-01],
            [-1.3060e-01],
            [-2.0210e-01],
            [-3.4880e-01],
            [-2.8876e-01],
            [-2.6767e-01],
            [-1.1433e-01],
            [-1.1392e-02],
            [ 1.0221e-03],
            [-9.4465e-04],
            [ 8.2190e-04],
            [ 1.3099e-05],
            [-2.8712e-04],
            [-4.3223e-04],
            [-1.2362e-01],
            [-2.1913e-01],
            [-2.4301e-01],
            [-3.1858e-01],
            [-1.8942e-01],
            [-4.9437e-01],
            [-4.0525e-01],
            [-6.8013e-02],
            [ 2.2953e-01],
            [ 7.2128e-01],
            [ 6.0160e-01],
            [ 2.8288e-01],
            [-2.3352e-01],
            [-5.7711e-01],
            [-1.5661e-01],
            [-1.4581e-01],
            [-3.6273e-01],
            [-2.7437e-01],
            [-2.5821e-01],
            [-2.0031e-01],
            [ 4.9633e-02],
            [ 4.1612e-02],
            [-1.8513e-03],
            [-8.5597e-04],
            [ 1.8215e-03],
            [ 6.1049e-04],
            [ 1.9938e-04],
            [-3.9272e-03],
            [-1.3940e-01],
            [-2.4405e-01],
            [-2.4828e-01],
            [-2.3565e-01],
            [-2.5946e-01],
            [-5.2036e-01],
            [-3.3285e-01],
            [-3.6621e-02],
            [ 2.4680e-01],
            [ 4.9221e-01],
            [ 2.3458e-01],
            [ 2.2097e-01],
            [-2.1094e-01],
            [-4.4383e-01],
            [-2.2330e-01],
            [-2.3315e-01],
            [-3.1203e-01],
            [-1.9487e-01],
            [-2.2718e-01],
            [-1.5206e-01],
            [ 1.2305e-01],
            [ 1.9312e-01],
            [ 7.9685e-04],
            [-9.4103e-04],
            [-1.9861e-04],
            [ 9.9772e-04],
            [ 1.2707e-03],
            [-1.6639e-03],
            [-1.2999e-01],
            [-2.8166e-01],
            [-2.2050e-01],
            [-1.2297e-01],
            [-3.0777e-01],
            [-3.6948e-01],
            [-1.5410e-01],
            [-1.4780e-01],
            [ 1.8346e-01],
            [ 2.9009e-01],
            [ 1.3905e-01],
            [ 1.0984e-01],
            [-3.1573e-02],
            [-3.4354e-01],
            [-9.0129e-02],
            [-2.9714e-01],
            [-1.4827e-01],
            [-1.3932e-01],
            [-1.6603e-01],
            [-6.3056e-02],
            [ 1.8612e-01],
            [ 2.0052e-01],
            [-2.7134e-05],
            [ 1.5348e-03],
            [-3.0186e-04],
            [ 4.0938e-04],
            [-6.7632e-04],
            [-3.2617e-03],
            [-1.1639e-01],
            [-2.8797e-01],
            [-1.1856e-01],
            [-7.3154e-03],
            [-2.7798e-01],
            [-1.1855e-01],
            [ 4.5018e-02],
            [-6.7436e-02],
            [-1.2961e-02],
            [ 1.3939e-01],
            [ 1.7378e-01],
            [ 1.5358e-01],
            [-2.6778e-02],
            [-9.0053e-02],
            [-5.1133e-03],
            [-1.1684e-01],
            [-8.9528e-02],
            [-6.1045e-02],
            [-4.2182e-02],
            [ 9.0279e-02],
            [ 2.2339e-01],
            [ 2.0156e-01],
            [ 2.2467e-05],
            [-1.2550e-03],
            [-7.9524e-04],
            [-1.4592e-03],
            [-2.4341e-03],
            [-1.2126e-04],
            [-5.5019e-02],
            [-1.0981e-01],
            [ 9.3553e-02],
            [-4.1696e-02],
            [-3.7313e-02],
            [ 1.1134e-01],
            [ 3.7415e-02],
            [-5.1016e-02],
            [-1.7630e-01],
            [-6.0320e-02],
            [ 2.1020e-01],
            [ 2.3959e-01],
            [ 6.7396e-02],
            [ 2.9125e-01],
            [ 3.2695e-01],
            [ 1.9023e-01],
            [ 2.1860e-01],
            [ 3.9042e-02],
            [ 1.0556e-01],
            [ 2.2804e-01],
            [ 2.3307e-01],
            [ 1.3450e-01],
            [-3.6564e-04],
            [-3.3880e-05],
            [ 9.8750e-04],
            [-1.6907e-04],
            [-3.9272e-04],
            [ 3.9938e-03],
            [ 3.7149e-02],
            [ 1.4923e-01],
            [ 2.9007e-01],
            [ 2.6650e-01],
            [ 2.2152e-01],
            [ 1.8445e-01],
            [-1.2630e-01],
            [-2.3512e-01],
            [-3.1205e-01],
            [-5.5628e-02],
            [ 1.4409e-01],
            [ 2.5980e-02],
            [ 1.4971e-01],
            [ 2.9404e-01],
            [ 3.1616e-01],
            [ 2.5143e-01],
            [ 4.2843e-01],
            [ 2.9584e-01],
            [ 1.3059e-01],
            [ 1.3588e-01],
            [ 1.9232e-01],
            [ 4.7314e-02],
            [-1.1801e-03],
            [-1.6411e-03],
            [-1.1983e-04],
            [ 6.6655e-05],
            [ 9.1697e-05],
            [ 2.8531e-03],
            [ 1.4067e-01],
            [ 2.3087e-01],
            [ 1.3467e-01],
            [ 9.4048e-02],
            [ 1.6608e-01],
            [ 2.4181e-01],
            [-1.5614e-01],
            [-3.5259e-01],
            [-4.4604e-01],
            [-2.4049e-02],
            [-9.3430e-02],
            [-1.5706e-01],
            [-4.9127e-02],
            [ 7.7699e-02],
            [ 1.2275e-01],
            [ 9.6132e-02],
            [ 3.0807e-01],
            [ 2.0825e-01],
            [ 2.7597e-02],
            [ 1.6552e-02],
            [ 2.5066e-02],
            [ 1.2060e-02],
            [-1.3505e-04],
            [ 1.7110e-06],
            [-2.3733e-04],
            [-2.2122e-04],
            [-1.4611e-03],
            [ 1.3784e-04],
            [ 4.9559e-02],
            [ 5.1187e-02],
            [-1.3546e-02],
            [-1.2889e-01],
            [-6.6621e-02],
            [ 2.3818e-01],
            [-7.1718e-02],
            [-2.7349e-01],
            [-1.9880e-01],
            [-6.2190e-02],
            [-2.0397e-01],
            [-7.2680e-02],
            [-1.6574e-05],
            [ 2.1952e-02],
            [ 3.4929e-02],
            [ 6.1684e-02],
            [ 2.0057e-01],
            [ 3.1400e-02],
            [ 7.8993e-03],
            [-6.6521e-04],
            [ 7.4936e-04],
            [-4.6069e-04],
            [-3.7327e-04],
            [-1.8720e-03],
            [ 1.5838e-03],
            [ 3.6597e-04],
            [-9.1368e-04],
            [-2.1601e-04],
            [-2.2191e-04],
            [-1.0846e-02],
            [-8.8628e-03],
            [-3.4458e-04],
            [-1.1214e-02],
            [ 8.3239e-02],
            [ 1.9510e-01],
            [ 4.6755e-02],
            [-5.2426e-02],
            [-2.4506e-02],
            [-1.6064e-01],
            [ 3.4662e-02],
            [ 6.1154e-02],
            [ 8.5607e-02],
            [ 2.0800e-02],
            [ 3.6094e-02],
            [ 7.1468e-03],
            [ 2.1759e-02],
            [ 1.1845e-02],
            [-2.4645e-04],
            [-2.9535e-04],
            [ 9.8739e-04],
            [-1.4327e-04],
            [ 2.1999e-04],
            [-1.3144e-03],
            [-1.7753e-03],
            [-7.5568e-04],
            [ 3.6761e-04],
            [ 6.2712e-05],
            [ 1.7757e-03],
            [ 2.0827e-03],
            [ 4.9462e-03],
            [ 1.0085e-02],
            [ 5.6281e-03],
            [ 2.9948e-03],
            [ 1.1400e-03],
            [ 1.9878e-01],
            [ 1.5982e-01],
            [ 5.8938e-02],
            [ 8.2638e-02],
            [ 1.6660e-02],
            [-1.4282e-03],
            [ 9.1127e-04],
            [ 9.2726e-04],
            [ 3.7004e-04],
            [ 6.4004e-04],
            [ 6.7870e-04],
            [ 1.2577e-03],
            [ 1.4228e-03],
            [-9.2635e-05],
            [-6.3314e-04],
            [-6.4067e-04],
            [ 2.4250e-04],
            [-9.5488e-04],
            [-4.8476e-04],
            [ 7.3312e-05],
            [-1.7000e-03],
            [ 5.0604e-04],
            [ 9.6993e-04],
            [ 1.3175e-03],
            [ 3.4707e-04],
            [-8.7314e-04],
            [ 2.7642e-04],
            [-8.8469e-04],
            [ 2.7667e-04],
            [ 9.8199e-04],
            [ 1.2911e-02],
            [ 3.9742e-02],
            [ 1.3516e-02],
            [ 3.3090e-04],
            [-2.1137e-04],
            [-3.5136e-04],
            [ 1.6309e-03],
            [ 1.0554e-03],
            [ 7.3709e-04],
            [ 9.0048e-04],
            [-7.0697e-04],
            [-3.1654e-04],
            [ 3.1712e-04],
            [ 4.6010e-04],
            [ 9.5045e-04],
            [-1.0009e-03],
            [-1.6685e-04],
            [ 1.4990e-03],
            [-4.2581e-04],
            [-2.0531e-03],
            [ 6.7953e-04],
            [ 4.8103e-04],
            [-2.5930e-04],
            [-3.6222e-04],
            [-3.9699e-04],
            [ 2.3928e-04],
            [-7.9614e-04],
            [-2.1639e-04],
            [ 9.0970e-04],
            [ 3.2931e-04],
            [-6.8675e-04],
            [ 1.7413e-04],
            [ 1.3325e-03],
            [-1.1063e-03],
            [ 3.5013e-04],
            [ 6.3391e-04],
            [ 4.8646e-04],
            [-9.5648e-04],
            [-4.4012e-04],
            [-1.3617e-04],
            [ 2.0291e-04],
            [-9.1338e-04]])
    [stdout:1]
    [DG] Rank 2: Global Gradient Norm: 0.00021530537924263626
    tensor([[-5.4105e-04],
            [-1.3045e-03],
            [-2.3855e-04],
            [ 5.8681e-04],
            [ 1.9476e-04],
            [ 8.4308e-04],
            [ 1.3606e-04],
            [-3.7884e-04],
            [-5.0034e-04],
            [-1.1751e-03],
            [-1.3683e-03],
            [-1.5652e-04],
            [-4.2702e-04],
            [-9.7871e-04],
            [-4.8231e-04],
            [ 7.7316e-04],
            [ 1.7644e-04],
            [-1.3649e-04],
            [-8.1466e-04],
            [ 6.9329e-04],
            [ 7.4197e-05],
            [ 3.2124e-04],
            [ 8.6270e-04],
            [-5.1548e-04],
            [-3.7360e-04],
            [-1.9439e-05],
            [-1.3404e-03],
            [ 4.8689e-04],
            [ 4.1458e-04],
            [-1.3369e-03],
            [-1.6877e-04],
            [ 4.8695e-04],
            [ 6.4009e-04],
            [-5.0863e-04],
            [ 1.7411e-03],
            [-5.0160e-04],
            [-7.7192e-04],
            [-4.0259e-04],
            [-6.6781e-04],
            [ 2.9919e-04],
            [-6.9413e-04],
            [ 5.2942e-04],
            [ 1.4289e-03],
            [-2.4318e-03],
            [-1.3059e-02],
            [-5.9160e-03],
            [-5.0573e-04],
            [-4.5606e-04],
            [-1.0758e-03],
            [ 6.8586e-04],
            [ 5.2861e-04],
            [-1.3040e-04],
            [-5.3096e-04],
            [ 7.1237e-04],
            [-8.9679e-04],
            [-8.7266e-04],
            [-4.1328e-05],
            [-1.7905e-03],
            [ 2.0443e-03],
            [-4.7453e-04],
            [-7.3630e-05],
            [ 1.6412e-03],
            [-1.5544e-04],
            [-1.1315e-03],
            [-3.4253e-04],
            [-3.8722e-04],
            [ 1.1653e-03],
            [-1.0911e-04],
            [ 7.2506e-03],
            [ 1.0557e-01],
            [ 3.3088e-02],
            [ 1.4403e-02],
            [ 6.2711e-02],
            [ 1.3494e-01],
            [ 3.2523e-02],
            [ 3.4457e-03],
            [ 9.5142e-04],
            [-3.7271e-04],
            [-1.4328e-03],
            [ 2.5682e-04],
            [-1.0477e-03],
            [ 9.1061e-04],
            [-2.0199e-05],
            [-5.4775e-04],
            [-1.6562e-03],
            [ 9.7708e-04],
            [ 1.2736e-04],
            [-7.6910e-04],
            [ 4.7750e-04],
            [ 6.0176e-04],
            [-1.1634e-03],
            [-9.7313e-04],
            [ 8.7060e-03],
            [ 1.1736e-02],
            [ 1.5709e-03],
            [ 2.3838e-03],
            [ 5.5873e-02],
            [ 1.1790e-01],
            [ 1.1921e-01],
            [ 1.2434e-01],
            [ 2.3453e-01],
            [ 2.7110e-01],
            [ 1.2929e-01],
            [ 1.7887e-02],
            [ 6.8535e-03],
            [ 2.5573e-03],
            [-1.3939e-03],
            [-1.6076e-03],
            [-4.8773e-04],
            [ 6.0169e-04],
            [ 1.9273e-03],
            [ 3.3049e-04],
            [ 7.1581e-04],
            [-9.0034e-04],
            [ 3.0410e-04],
            [ 1.4058e-04],
            [-8.2944e-04],
            [ 1.5920e-03],
            [-2.2250e-04],
            [-5.3122e-03],
            [ 8.4993e-02],
            [ 1.2974e-01],
            [ 4.0902e-02],
            [ 5.0917e-02],
            [ 1.1401e-01],
            [ 1.2502e-01],
            [ 5.3357e-02],
            [ 1.2477e-01],
            [ 7.3475e-02],
            [ 2.0133e-01],
            [-9.1516e-02],
            [ 3.6535e-02],
            [ 1.5575e-01],
            [ 9.9466e-02],
            [ 7.6248e-02],
            [ 4.0132e-02],
            [ 3.8402e-03],
            [-4.6554e-04],
            [-1.3029e-03],
            [ 1.5080e-04],
            [-1.0466e-03],
            [-1.5949e-04],
            [-2.4158e-04],
            [-6.3473e-04],
            [-1.4772e-03],
            [ 6.2617e-03],
            [-2.3464e-03],
            [-6.8381e-03],
            [ 1.2517e-01],
            [ 1.4357e-01],
            [ 3.1361e-02],
            [ 5.0578e-02],
            [ 1.7412e-01],
            [ 2.8534e-01],
            [-4.7071e-02],
            [-1.4693e-01],
            [-2.0356e-04],
            [ 6.8268e-02],
            [-1.4395e-01],
            [-5.6544e-02],
            [ 2.1238e-01],
            [ 1.8764e-01],
            [ 1.5097e-01],
            [ 1.3172e-01],
            [ 5.1673e-02],
            [ 5.6072e-04],
            [ 2.5057e-04],
            [-1.7655e-03],
            [ 1.7811e-03],
            [-1.7870e-04],
            [ 1.3582e-04],
            [ 9.0135e-04],
            [ 2.0179e-03],
            [ 2.2173e-03],
            [ 1.8398e-02],
            [ 7.3196e-02],
            [ 7.9233e-02],
            [ 2.6257e-02],
            [-4.0096e-02],
            [-6.2816e-02],
            [ 7.4902e-02],
            [ 1.1035e-01],
            [-4.8329e-02],
            [-1.1968e-01],
            [-1.4363e-01],
            [-6.7106e-02],
            [-5.6265e-02],
            [-1.1003e-01],
            [ 2.0983e-01],
            [ 2.1829e-01],
            [ 9.8755e-03],
            [ 3.3544e-02],
            [ 4.4741e-02],
            [-8.3623e-04],
            [-8.5949e-05],
            [-2.6526e-04],
            [-1.1790e-04],
            [ 1.4127e-04],
            [-5.1918e-04],
            [ 7.9258e-04],
            [ 1.3269e-01],
            [ 1.3985e-01],
            [ 6.2427e-02],
            [ 1.4624e-01],
            [ 1.7466e-02],
            [-1.5585e-01],
            [-2.4965e-01],
            [-2.8416e-01],
            [-4.6471e-02],
            [ 4.4980e-02],
            [-7.3328e-03],
            [-2.9527e-02],
            [ 1.0888e-02],
            [-1.7594e-01],
            [-1.3119e-01],
            [-1.4207e-01],
            [ 5.6166e-02],
            [-5.2570e-02],
            [-1.0742e-01],
            [-5.8749e-02],
            [-6.2239e-02],
            [-1.2640e-02],
            [-1.4508e-05],
            [-3.1371e-04],
            [-8.2600e-04],
            [-3.4794e-04],
            [-1.4636e-04],
            [-1.1128e-03],
            [ 1.2306e-01],
            [ 1.6096e-01],
            [ 6.1389e-02],
            [ 1.1082e-01],
            [ 1.1766e-02],
            [-1.6376e-01],
            [-2.0334e-01],
            [-1.4668e-01],
            [ 1.4204e-01],
            [ 3.8477e-03],
            [ 2.4888e-01],
            [ 1.5759e-01],
            [-2.0614e-01],
            [-2.5094e-01],
            [ 7.2122e-04],
            [-1.6129e-02],
            [-1.9733e-01],
            [-2.7025e-01],
            [-1.5386e-01],
            [-6.2906e-02],
            [-3.4288e-02],
            [-4.0518e-03],
            [-7.1764e-04],
            [ 6.6908e-04],
            [ 3.2718e-04],
            [ 3.0669e-04],
            [-1.1061e-04],
            [ 1.2584e-04],
            [ 1.4936e-01],
            [ 1.3547e-01],
            [-4.3119e-02],
            [ 4.7758e-02],
            [ 1.9646e-02],
            [-3.9201e-02],
            [-8.7595e-02],
            [-8.8336e-03],
            [ 1.2637e-01],
            [ 2.7384e-01],
            [ 2.9850e-01],
            [ 1.1623e-01],
            [-1.7446e-02],
            [-9.3210e-02],
            [ 1.6622e-01],
            [ 1.2816e-01],
            [-2.8506e-01],
            [-3.9594e-01],
            [-2.4441e-01],
            [-1.3550e-01],
            [-2.4789e-02],
            [-5.4242e-03],
            [ 1.0709e-03],
            [ 7.1802e-04],
            [-1.4283e-03],
            [ 2.1113e-04],
            [ 6.1421e-04],
            [-1.3334e-03],
            [ 1.4552e-01],
            [ 3.9650e-02],
            [-1.2435e-01],
            [-1.2060e-01],
            [-1.1163e-01],
            [-1.8812e-02],
            [-3.6345e-02],
            [ 4.3420e-02],
            [ 1.0976e-01],
            [ 1.1508e-01],
            [ 1.1580e-01],
            [ 2.6903e-01],
            [ 2.1422e-01],
            [-1.0324e-01],
            [ 1.3432e-02],
            [-2.2741e-02],
            [-3.0328e-01],
            [-3.4095e-01],
            [-2.8226e-01],
            [-1.6030e-01],
            [-4.6170e-02],
            [-5.3130e-03],
            [ 4.1101e-04],
            [ 6.9317e-04],
            [-9.8351e-04],
            [ 6.0529e-04],
            [-2.1607e-04],
            [ 6.4386e-04],
            [ 1.3726e-01],
            [-3.8051e-02],
            [-1.6696e-01],
            [-1.8973e-01],
            [-2.2579e-01],
            [-1.6501e-01],
            [-1.2378e-01],
            [-9.3703e-02],
            [-8.1240e-03],
            [ 4.7210e-02],
            [ 1.2682e-01],
            [ 5.0609e-01],
            [ 3.3579e-01],
            [-2.5666e-02],
            [-1.7571e-01],
            [-2.5668e-01],
            [-1.2065e-01],
            [-2.5367e-01],
            [-2.5888e-01],
            [-8.9716e-02],
            [-7.8603e-02],
            [-6.2367e-03],
            [-3.7541e-04],
            [ 1.5716e-03],
            [-1.7036e-04],
            [ 3.6778e-04],
            [-1.4130e-03],
            [ 1.2263e-04],
            [ 7.5421e-02],
            [-1.2736e-01],
            [-1.7615e-01],
            [-1.6302e-01],
            [-2.7537e-01],
            [-2.9733e-01],
            [-1.2279e-01],
            [-1.3726e-01],
            [-2.3564e-01],
            [ 1.6874e-01],
            [ 3.7477e-01],
            [ 6.8394e-01],
            [ 3.1668e-01],
            [-7.4922e-03],
            [-3.2702e-01],
            [-3.2107e-01],
            [-7.4969e-02],
            [-2.9166e-01],
            [-2.6684e-01],
            [-1.7898e-01],
            [-9.3595e-02],
            [-9.3452e-03],
            [ 1.7086e-03],
            [ 1.0012e-03],
            [-1.3916e-04],
            [-3.2752e-04],
            [ 9.5726e-04],
            [-3.7799e-04],
            [ 2.6437e-03],
            [-1.6818e-01],
            [-1.7428e-01],
            [-2.7102e-01],
            [-3.4440e-01],
            [-2.2166e-01],
            [-1.7838e-01],
            [-3.4935e-01],
            [-1.6612e-01],
            [ 4.6741e-01],
            [ 7.1802e-01],
            [ 6.3802e-01],
            [ 2.7221e-01],
            [-7.5658e-02],
            [-6.6731e-01],
            [-2.4221e-01],
            [ 2.0253e-04],
            [-4.0784e-01],
            [-4.2215e-01],
            [-2.6032e-01],
            [-1.0320e-01],
            [-1.4749e-02],
            [ 3.5213e-04],
            [ 1.3588e-03],
            [-2.0972e-03],
            [ 5.5523e-04],
            [-9.4352e-06],
            [-1.7758e-03],
            [-3.8491e-02],
            [-1.8131e-01],
            [-2.5295e-01],
            [-3.6540e-01],
            [-2.5232e-01],
            [-6.6879e-02],
            [-4.1002e-01],
            [-3.1494e-01],
            [-2.2827e-02],
            [ 5.3574e-01],
            [ 8.9542e-01],
            [ 5.5001e-01],
            [ 2.9052e-01],
            [-2.5619e-01],
            [-8.2296e-01],
            [-1.4310e-01],
            [-1.4138e-01],
            [-4.0839e-01],
            [-3.4570e-01],
            [-2.6477e-01],
            [-1.1727e-01],
            [-1.4586e-02],
            [ 1.3353e-03],
            [ 6.7560e-04],
            [ 1.0227e-03],
            [ 8.2836e-04],
            [ 8.9151e-05],
            [ 2.4532e-03],
            [-1.1113e-01],
            [-2.2878e-01],
            [-2.8006e-01],
            [-2.9946e-01],
            [-2.1536e-01],
            [-1.9874e-01],
            [-4.9707e-01],
            [-2.4265e-01],
            [ 2.0239e-01],
            [ 7.5617e-01],
            [ 9.3105e-01],
            [ 3.3986e-01],
            [ 3.1344e-02],
            [-5.7400e-01],
            [-4.9098e-01],
            [-1.3145e-01],
            [-2.0286e-01],
            [-3.4878e-01],
            [-2.8986e-01],
            [-2.6787e-01],
            [-1.1292e-01],
            [-1.2163e-02],
            [-1.1126e-03],
            [ 4.8481e-04],
            [-1.4076e-03],
            [-2.6372e-05],
            [ 2.6181e-04],
            [-1.4498e-03],
            [-1.2360e-01],
            [-2.2155e-01],
            [-2.4261e-01],
            [-3.2201e-01],
            [-1.9169e-01],
            [-4.9383e-01],
            [-4.0304e-01],
            [-6.7936e-02],
            [ 2.2788e-01],
            [ 7.2298e-01],
            [ 6.0458e-01],
            [ 2.8036e-01],
            [-2.3338e-01],
            [-5.7792e-01],
            [-1.5691e-01],
            [-1.4801e-01],
            [-3.6394e-01],
            [-2.7390e-01],
            [-2.5580e-01],
            [-1.9890e-01],
            [ 4.6839e-02],
            [ 4.0523e-02],
            [-3.0278e-04],
            [-1.4856e-05],
            [-1.2999e-03],
            [-6.0532e-04],
            [-9.5009e-04],
            [-1.6366e-04],
            [-1.3928e-01],
            [-2.4609e-01],
            [-2.4982e-01],
            [-2.3467e-01],
            [-2.5646e-01],
            [-5.1998e-01],
            [-3.3092e-01],
            [-3.5753e-02],
            [ 2.4536e-01],
            [ 4.9198e-01],
            [ 2.3509e-01],
            [ 2.2028e-01],
            [-2.0988e-01],
            [-4.4443e-01],
            [-2.2430e-01],
            [-2.3105e-01],
            [-3.1525e-01],
            [-1.9629e-01],
            [-2.2810e-01],
            [-1.5301e-01],
            [ 1.2347e-01],
            [ 1.9081e-01],
            [ 4.4079e-04],
            [ 1.7716e-04],
            [ 3.2132e-04],
            [-7.7858e-04],
            [ 2.3815e-04],
            [-7.8859e-04],
            [-1.3022e-01],
            [-2.8199e-01],
            [-2.2169e-01],
            [-1.2326e-01],
            [-3.0509e-01],
            [-3.6700e-01],
            [-1.5383e-01],
            [-1.4906e-01],
            [ 1.8345e-01],
            [ 2.9216e-01],
            [ 1.3669e-01],
            [ 1.1253e-01],
            [-3.0831e-02],
            [-3.4478e-01],
            [-8.9930e-02],
            [-2.9823e-01],
            [-1.4984e-01],
            [-1.3869e-01],
            [-1.6807e-01],
            [-6.2948e-02],
            [ 1.8717e-01],
            [ 1.9890e-01],
            [-1.6798e-04],
            [-4.0555e-04],
            [ 2.7193e-04],
            [-1.5775e-03],
            [-1.5088e-04],
            [-2.1853e-03],
            [-1.1728e-01],
            [-2.8792e-01],
            [-1.2023e-01],
            [-5.1142e-03],
            [-2.7694e-01],
            [-1.1901e-01],
            [ 4.2591e-02],
            [-6.9015e-02],
            [-1.5056e-02],
            [ 1.3946e-01],
            [ 1.7302e-01],
            [ 1.5166e-01],
            [-2.6023e-02],
            [-9.0450e-02],
            [-4.4404e-03],
            [-1.1735e-01],
            [-8.8222e-02],
            [-6.0719e-02],
            [-4.2707e-02],
            [ 8.9974e-02],
            [ 2.2595e-01],
            [ 1.9918e-01],
            [ 1.2185e-04],
            [-9.0256e-05],
            [ 5.4604e-04],
            [-1.4244e-03],
            [-6.9972e-04],
            [-8.3422e-04],
            [-5.7338e-02],
            [-1.0847e-01],
            [ 9.2992e-02],
            [-4.1183e-02],
            [-3.5627e-02],
            [ 1.1282e-01],
            [ 3.9849e-02],
            [-5.2756e-02],
            [-1.7426e-01],
            [-6.1121e-02],
            [ 2.1046e-01],
            [ 2.3982e-01],
            [ 6.7214e-02],
            [ 2.9174e-01],
            [ 3.2445e-01],
            [ 1.8825e-01],
            [ 2.1916e-01],
            [ 4.0580e-02],
            [ 1.0622e-01],
            [ 2.2637e-01],
            [ 2.3247e-01],
            [ 1.3515e-01],
            [ 4.7845e-04],
            [-1.3829e-03],
            [ 4.1264e-04],
            [-1.3693e-03],
            [ 5.5390e-04],
            [ 2.3882e-03],
            [ 3.8778e-02],
            [ 1.4962e-01],
            [ 2.9471e-01],
            [ 2.6574e-01],
            [ 2.2289e-01],
            [ 1.8516e-01],
            [-1.2570e-01],
            [-2.3436e-01],
            [-3.1220e-01],
            [-5.6212e-02],
            [ 1.4425e-01],
            [ 2.5748e-02],
            [ 1.5027e-01],
            [ 2.9327e-01],
            [ 3.1635e-01],
            [ 2.5280e-01],
            [ 4.2678e-01],
            [ 2.9693e-01],
            [ 1.3044e-01],
            [ 1.3873e-01],
            [ 1.9174e-01],
            [ 4.6516e-02],
            [-1.9952e-03],
            [-3.7283e-05],
            [-6.0973e-04],
            [ 1.1016e-04],
            [ 3.1679e-04],
            [ 3.8434e-03],
            [ 1.3903e-01],
            [ 2.3070e-01],
            [ 1.3292e-01],
            [ 9.6507e-02],
            [ 1.6755e-01],
            [ 2.4343e-01],
            [-1.5567e-01],
            [-3.5226e-01],
            [-4.4767e-01],
            [-2.0905e-02],
            [-9.2010e-02],
            [-1.5593e-01],
            [-4.9366e-02],
            [ 7.6768e-02],
            [ 1.2089e-01],
            [ 9.7733e-02],
            [ 3.0946e-01],
            [ 2.0932e-01],
            [ 2.7068e-02],
            [ 1.7411e-02],
            [ 2.6624e-02],
            [ 8.1774e-03],
            [-8.4713e-04],
            [-1.0216e-03],
            [ 6.7928e-04],
            [ 1.6038e-04],
            [ 7.8992e-04],
            [ 1.1912e-04],
            [ 4.7260e-02],
            [ 4.7792e-02],
            [-1.1222e-02],
            [-1.2904e-01],
            [-6.3993e-02],
            [ 2.3958e-01],
            [-6.9874e-02],
            [-2.7389e-01],
            [-1.9822e-01],
            [-6.0978e-02],
            [-2.0312e-01],
            [-7.1214e-02],
            [ 2.2404e-03],
            [ 2.3757e-02],
            [ 3.3307e-02],
            [ 6.2230e-02],
            [ 1.9970e-01],
            [ 3.0685e-02],
            [ 7.5767e-03],
            [-8.1805e-04],
            [ 1.8659e-04],
            [-1.8703e-04],
            [-2.3103e-04],
            [ 4.7339e-04],
            [-3.5506e-04],
            [-1.2826e-03],
            [-2.3226e-04],
            [ 1.1795e-04],
            [-1.3997e-03],
            [-9.4841e-03],
            [-6.8245e-03],
            [-1.8802e-03],
            [-8.1506e-03],
            [ 8.4905e-02],
            [ 1.9244e-01],
            [ 5.0709e-02],
            [-5.2487e-02],
            [-2.4234e-02],
            [-1.6134e-01],
            [ 3.2152e-02],
            [ 6.0570e-02],
            [ 8.5733e-02],
            [ 2.0480e-02],
            [ 3.7413e-02],
            [ 7.9522e-03],
            [ 1.9795e-02],
            [ 1.0564e-02],
            [-4.3836e-04],
            [ 1.4202e-03],
            [ 3.5241e-04],
            [-1.0990e-04],
            [ 1.1342e-03],
            [-7.9986e-04],
            [-1.9340e-04],
            [ 1.9058e-03],
            [-1.4058e-03],
            [ 3.6241e-04],
            [-1.1811e-03],
            [-1.1765e-04],
            [ 4.1137e-03],
            [ 9.7826e-03],
            [ 4.8273e-03],
            [ 6.4678e-04],
            [ 1.4480e-03],
            [ 2.0287e-01],
            [ 1.5944e-01],
            [ 6.0246e-02],
            [ 8.3676e-02],
            [ 2.0088e-02],
            [ 4.0564e-04],
            [-4.2532e-04],
            [-1.1658e-03],
            [ 1.3472e-03],
            [-1.4504e-03],
            [-8.1858e-06],
            [-1.4321e-03],
            [-1.2969e-04],
            [-1.5210e-04],
            [-1.9471e-04],
            [-6.4982e-04],
            [-3.7121e-04],
            [-1.4698e-03],
            [-2.7322e-04],
            [ 2.6307e-04],
            [ 6.5332e-04],
            [ 6.6191e-04],
            [-2.8233e-05],
            [ 1.6553e-04],
            [-4.1361e-04],
            [ 1.1989e-03],
            [ 4.1916e-04],
            [-7.6089e-04],
            [-3.1575e-04],
            [-5.6632e-04],
            [ 1.3469e-02],
            [ 4.0405e-02],
            [ 1.2751e-02],
            [-9.2466e-04],
            [ 2.4969e-04],
            [ 5.6915e-04],
            [ 7.2845e-04],
            [-1.2454e-03],
            [ 5.8202e-04],
            [-3.4049e-04],
            [ 2.8042e-03],
            [ 5.4605e-04],
            [ 2.3954e-04],
            [-7.7665e-04],
            [ 6.5562e-04],
            [-4.6199e-04],
            [ 2.3023e-03],
            [ 2.6098e-05],
            [-8.5694e-05],
            [ 1.0787e-03],
            [ 3.9381e-04],
            [-9.1379e-05],
            [ 4.9962e-04],
            [-2.3306e-05],
            [-2.2711e-04],
            [ 1.2350e-03],
            [ 1.3516e-03],
            [ 8.7773e-04],
            [ 8.8540e-05],
            [ 1.0371e-03],
            [ 2.6188e-04],
            [-9.2860e-04],
            [-1.7155e-03],
            [ 1.3991e-03],
            [-3.7069e-04],
            [ 1.1849e-03],
            [ 2.9149e-04],
            [ 4.5626e-04],
            [-7.9581e-04],
            [ 3.9340e-04],
            [-1.0818e-03],
            [-2.0235e-04]])
    [stdout:2]
    [DG] Rank 1: Global Gradient Norm: 0.00021530537924263626
    tensor([[-1.4601e-03],
            [ 3.6807e-04],
            [ 2.8047e-04],
            [-1.1137e-03],
            [-5.0861e-04],
            [-6.3935e-04],
            [ 1.2964e-03],
            [ 1.6874e-06],
            [ 7.8016e-04],
            [ 1.4839e-03],
            [ 7.2005e-04],
            [ 1.2088e-03],
            [ 1.9444e-04],
            [-2.8630e-04],
            [-1.2475e-04],
            [ 4.5440e-04],
            [-1.1125e-03],
            [ 1.6875e-05],
            [-1.4399e-04],
            [ 2.1096e-03],
            [ 5.9366e-04],
            [-7.3517e-04],
            [-1.4586e-03],
            [ 5.4612e-04],
            [-2.0658e-03],
            [ 1.3729e-04],
            [ 1.5754e-04],
            [-8.8431e-04],
            [-5.9968e-04],
            [-2.0020e-04],
            [ 1.8262e-04],
            [-4.4049e-04],
            [-4.0506e-04],
            [ 1.0968e-03],
            [-5.4665e-04],
            [-8.3940e-04],
            [-2.4617e-04],
            [ 3.1073e-04],
            [ 5.7859e-04],
            [ 7.2810e-04],
            [ 4.4404e-04],
            [-4.9193e-04],
            [ 5.2091e-04],
            [-4.3049e-03],
            [-1.2800e-02],
            [-7.2577e-03],
            [-1.4582e-03],
            [-1.1099e-03],
            [-7.5963e-04],
            [ 1.0710e-03],
            [ 7.6558e-04],
            [-2.8874e-04],
            [ 1.6803e-04],
            [-2.4144e-03],
            [ 1.0284e-05],
            [ 4.3448e-04],
            [-3.3786e-04],
            [ 4.6066e-04],
            [ 1.1825e-04],
            [-6.2553e-04],
            [-2.6284e-04],
            [-2.4203e-03],
            [ 5.0401e-04],
            [-1.5230e-03],
            [-1.2865e-03],
            [ 2.5686e-04],
            [-7.6880e-04],
            [ 1.6852e-03],
            [ 7.1085e-03],
            [ 1.0539e-01],
            [ 3.2410e-02],
            [ 1.4540e-02],
            [ 6.1953e-02],
            [ 1.3626e-01],
            [ 3.2894e-02],
            [ 4.0087e-03],
            [ 1.8477e-03],
            [-9.4369e-04],
            [-2.0496e-04],
            [ 2.9792e-04],
            [ 2.0598e-04],
            [-3.6417e-04],
            [ 4.6086e-04],
            [-4.5173e-04],
            [-1.5655e-03],
            [-5.2483e-04],
            [ 1.7115e-03],
            [-3.8123e-04],
            [ 8.0392e-04],
            [-4.6637e-04],
            [-1.5653e-03],
            [-6.6233e-04],
            [ 8.8392e-03],
            [ 1.1146e-02],
            [ 3.6789e-03],
            [ 1.2208e-04],
            [ 5.6982e-02],
            [ 1.1654e-01],
            [ 1.1959e-01],
            [ 1.2524e-01],
            [ 2.3485e-01],
            [ 2.7053e-01],
            [ 1.3099e-01],
            [ 1.9201e-02],
            [ 8.0157e-03],
            [ 5.0037e-03],
            [-1.6080e-04],
            [ 1.4427e-04],
            [-4.1520e-04],
            [-2.6703e-04],
            [ 5.5270e-04],
            [-4.0053e-04],
            [-1.6603e-04],
            [ 1.7869e-04],
            [-7.0400e-04],
            [ 1.8997e-03],
            [ 1.1895e-03],
            [ 2.7376e-03],
            [-1.0166e-03],
            [-5.9550e-03],
            [ 8.6449e-02],
            [ 1.3113e-01],
            [ 4.4642e-02],
            [ 5.0200e-02],
            [ 1.1273e-01],
            [ 1.2555e-01],
            [ 5.5697e-02],
            [ 1.2269e-01],
            [ 7.3302e-02],
            [ 2.0096e-01],
            [-9.2668e-02],
            [ 3.5513e-02],
            [ 1.5682e-01],
            [ 1.0262e-01],
            [ 7.7001e-02],
            [ 3.8689e-02],
            [ 4.9422e-03],
            [-6.7594e-04],
            [ 2.7102e-03],
            [-1.1427e-03],
            [ 6.7161e-04],
            [-7.0064e-04],
            [ 4.2528e-04],
            [-6.1524e-04],
            [ 1.2401e-03],
            [ 7.0024e-03],
            [-9.1142e-04],
            [-7.5670e-03],
            [ 1.2558e-01],
            [ 1.4368e-01],
            [ 3.2798e-02],
            [ 4.8637e-02],
            [ 1.7321e-01],
            [ 2.8823e-01],
            [-4.5876e-02],
            [-1.4739e-01],
            [ 6.1990e-04],
            [ 6.8185e-02],
            [-1.4369e-01],
            [-5.7728e-02],
            [ 2.0998e-01],
            [ 1.8860e-01],
            [ 1.5111e-01],
            [ 1.2975e-01],
            [ 5.3541e-02],
            [-1.8383e-04],
            [-3.6953e-04],
            [-5.1051e-04],
            [-2.7961e-04],
            [-4.3549e-04],
            [-1.0275e-03],
            [ 5.0856e-04],
            [ 7.4130e-04],
            [ 8.1603e-04],
            [ 1.7658e-02],
            [ 7.5267e-02],
            [ 8.0049e-02],
            [ 2.5303e-02],
            [-4.0413e-02],
            [-6.3898e-02],
            [ 7.5162e-02],
            [ 1.0742e-01],
            [-5.0474e-02],
            [-1.2061e-01],
            [-1.4516e-01],
            [-6.6619e-02],
            [-5.4875e-02],
            [-1.0849e-01],
            [ 2.0973e-01],
            [ 2.1737e-01],
            [ 9.4405e-03],
            [ 3.4244e-02],
            [ 4.4887e-02],
            [ 3.5142e-04],
            [ 3.5879e-04],
            [ 1.7530e-03],
            [-3.9285e-04],
            [ 1.3252e-03],
            [ 8.2911e-05],
            [-3.0471e-04],
            [ 1.3281e-01],
            [ 1.4017e-01],
            [ 6.1342e-02],
            [ 1.4695e-01],
            [ 1.9289e-02],
            [-1.5785e-01],
            [-2.4664e-01],
            [-2.8337e-01],
            [-4.4670e-02],
            [ 4.6129e-02],
            [-6.7574e-03],
            [-3.1061e-02],
            [ 1.1531e-02],
            [-1.7456e-01],
            [-1.3049e-01],
            [-1.3898e-01],
            [ 5.4945e-02],
            [-5.4815e-02],
            [-1.0925e-01],
            [-5.6284e-02],
            [-6.1578e-02],
            [-1.2639e-02],
            [ 1.5975e-03],
            [-6.8186e-04],
            [-7.5997e-04],
            [ 7.6285e-04],
            [ 4.0989e-05],
            [-2.4059e-03],
            [ 1.2357e-01],
            [ 1.6106e-01],
            [ 5.8476e-02],
            [ 1.0875e-01],
            [ 1.1463e-02],
            [-1.6338e-01],
            [-2.0518e-01],
            [-1.4605e-01],
            [ 1.4186e-01],
            [ 3.3907e-03],
            [ 2.4871e-01],
            [ 1.5548e-01],
            [-2.0760e-01],
            [-2.5226e-01],
            [ 5.9675e-04],
            [-1.6546e-02],
            [-1.9869e-01],
            [-2.7242e-01],
            [-1.5334e-01],
            [-6.3316e-02],
            [-3.5325e-02],
            [-4.5653e-03],
            [ 1.1924e-03],
            [-8.6510e-04],
            [ 8.1072e-04],
            [-1.5097e-03],
            [-2.9425e-04],
            [-1.0574e-04],
            [ 1.4686e-01],
            [ 1.3667e-01],
            [-4.0417e-02],
            [ 4.9615e-02],
            [ 1.9737e-02],
            [-3.7048e-02],
            [-8.7116e-02],
            [-6.6421e-03],
            [ 1.2685e-01],
            [ 2.7169e-01],
            [ 2.9841e-01],
            [ 1.1425e-01],
            [-1.4206e-02],
            [-9.2706e-02],
            [ 1.6673e-01],
            [ 1.2864e-01],
            [-2.8277e-01],
            [-3.9892e-01],
            [-2.4373e-01],
            [-1.3681e-01],
            [-2.6714e-02],
            [-5.1068e-03],
            [-8.3992e-04],
            [ 2.2701e-03],
            [ 2.0884e-04],
            [-4.6784e-04],
            [ 5.8062e-04],
            [-2.6994e-04],
            [ 1.4624e-01],
            [ 3.9924e-02],
            [-1.2122e-01],
            [-1.1986e-01],
            [-1.1219e-01],
            [-1.7419e-02],
            [-3.6747e-02],
            [ 4.3077e-02],
            [ 1.0947e-01],
            [ 1.1385e-01],
            [ 1.1517e-01],
            [ 2.6844e-01],
            [ 2.1590e-01],
            [-1.0568e-01],
            [ 1.4770e-02],
            [-2.2394e-02],
            [-3.0350e-01],
            [-3.4039e-01],
            [-2.8394e-01],
            [-1.5826e-01],
            [-4.6630e-02],
            [-5.6043e-03],
            [-1.0634e-03],
            [-1.2603e-03],
            [-3.6343e-04],
            [-6.5486e-04],
            [-3.2603e-05],
            [ 2.5604e-04],
            [ 1.3834e-01],
            [-3.7124e-02],
            [-1.6509e-01],
            [-1.9083e-01],
            [-2.2282e-01],
            [-1.6486e-01],
            [-1.2034e-01],
            [-9.3942e-02],
            [-8.4512e-03],
            [ 4.8999e-02],
            [ 1.2564e-01],
            [ 5.0363e-01],
            [ 3.3433e-01],
            [-2.4532e-02],
            [-1.7542e-01],
            [-2.5630e-01],
            [-1.2059e-01],
            [-2.5547e-01],
            [-2.5859e-01],
            [-9.1239e-02],
            [-7.6006e-02],
            [-7.5642e-03],
            [-2.4966e-04],
            [-1.3289e-04],
            [ 8.7858e-04],
            [-1.1220e-04],
            [ 1.3095e-05],
            [ 1.3748e-03],
            [ 7.5181e-02],
            [-1.2726e-01],
            [-1.7514e-01],
            [-1.6376e-01],
            [-2.7680e-01],
            [-2.9637e-01],
            [-1.2090e-01],
            [-1.3616e-01],
            [-2.3595e-01],
            [ 1.6756e-01],
            [ 3.7632e-01],
            [ 6.8313e-01],
            [ 3.1708e-01],
            [-6.4137e-03],
            [-3.2755e-01],
            [-3.2144e-01],
            [-7.3498e-02],
            [-2.9087e-01],
            [-2.6746e-01],
            [-1.7745e-01],
            [-9.4720e-02],
            [-9.1437e-03],
            [-1.0995e-03],
            [-1.1985e-03],
            [-2.6712e-04],
            [-1.3973e-03],
            [-2.1326e-03],
            [-4.1598e-04],
            [ 5.1409e-03],
            [-1.6729e-01],
            [-1.6994e-01],
            [-2.7094e-01],
            [-3.4222e-01],
            [-2.2267e-01],
            [-1.7828e-01],
            [-3.5014e-01],
            [-1.6620e-01],
            [ 4.6738e-01],
            [ 7.1557e-01],
            [ 6.3588e-01],
            [ 2.7188e-01],
            [-7.6461e-02],
            [-6.6673e-01],
            [-2.4187e-01],
            [ 4.3329e-04],
            [-4.0664e-01],
            [-4.2211e-01],
            [-2.6162e-01],
            [-1.0481e-01],
            [-1.4106e-02],
            [-4.4096e-04],
            [-2.4926e-04],
            [ 3.4233e-04],
            [-2.2094e-04],
            [ 1.2712e-04],
            [ 1.4019e-03],
            [-3.9597e-02],
            [-1.8035e-01],
            [-2.5124e-01],
            [-3.6599e-01],
            [-2.5217e-01],
            [-6.6799e-02],
            [-4.0960e-01],
            [-3.1429e-01],
            [-2.0585e-02],
            [ 5.3864e-01],
            [ 8.9614e-01],
            [ 5.5150e-01],
            [ 2.9108e-01],
            [-2.5698e-01],
            [-8.2434e-01],
            [-1.4394e-01],
            [-1.4168e-01],
            [-4.0969e-01],
            [-3.4564e-01],
            [-2.6448e-01],
            [-1.1539e-01],
            [-1.5720e-02],
            [ 3.8095e-05],
            [-1.2083e-03],
            [ 1.5459e-04],
            [ 1.6404e-04],
            [-2.8264e-04],
            [ 2.7905e-03],
            [-1.1089e-01],
            [-2.3128e-01],
            [-2.7966e-01],
            [-2.9946e-01],
            [-2.1618e-01],
            [-1.9684e-01],
            [-4.9770e-01],
            [-2.4385e-01],
            [ 2.0384e-01],
            [ 7.5631e-01],
            [ 9.3139e-01],
            [ 3.4109e-01],
            [ 3.1030e-02],
            [-5.7288e-01],
            [-4.9035e-01],
            [-1.3247e-01],
            [-2.0333e-01],
            [-3.4728e-01],
            [-2.8828e-01],
            [-2.6963e-01],
            [-1.1156e-01],
            [-9.8511e-03],
            [ 7.9592e-04],
            [-1.3185e-03],
            [-1.7262e-03],
            [ 5.3847e-04],
            [ 3.1418e-04],
            [-2.5126e-03],
            [-1.2475e-01],
            [-2.1821e-01],
            [-2.4315e-01],
            [-3.2002e-01],
            [-1.8809e-01],
            [-4.9593e-01],
            [-4.0567e-01],
            [-6.8002e-02],
            [ 2.2881e-01],
            [ 7.2002e-01],
            [ 6.0306e-01],
            [ 2.8115e-01],
            [-2.3334e-01],
            [-5.7732e-01],
            [-1.5680e-01],
            [-1.4443e-01],
            [-3.6374e-01],
            [-2.7462e-01],
            [-2.5875e-01],
            [-2.0191e-01],
            [ 4.8381e-02],
            [ 4.1172e-02],
            [ 2.5255e-03],
            [-1.2052e-04],
            [ 1.6646e-03],
            [-1.5279e-03],
            [-1.1358e-03],
            [-5.9405e-04],
            [-1.3847e-01],
            [-2.4521e-01],
            [-2.4884e-01],
            [-2.3339e-01],
            [-2.5936e-01],
            [-5.2167e-01],
            [-3.3141e-01],
            [-3.5674e-02],
            [ 2.4681e-01],
            [ 4.9201e-01],
            [ 2.3351e-01],
            [ 2.2056e-01],
            [-2.1137e-01],
            [-4.4494e-01],
            [-2.2043e-01],
            [-2.3310e-01],
            [-3.1326e-01],
            [-1.9369e-01],
            [-2.2855e-01],
            [-1.5262e-01],
            [ 1.2347e-01],
            [ 1.9013e-01],
            [ 7.4477e-04],
            [ 5.4482e-04],
            [-1.3418e-03],
            [ 4.5893e-04],
            [ 8.2499e-04],
            [-1.2459e-03],
            [-1.2725e-01],
            [-2.7959e-01],
            [-2.2025e-01],
            [-1.2268e-01],
            [-3.0743e-01],
            [-3.6950e-01],
            [-1.5442e-01],
            [-1.4748e-01],
            [ 1.8404e-01],
            [ 2.9180e-01],
            [ 1.3488e-01],
            [ 1.1156e-01],
            [-3.1117e-02],
            [-3.4438e-01],
            [-8.8033e-02],
            [-2.9613e-01],
            [-1.4755e-01],
            [-1.3955e-01],
            [-1.6737e-01],
            [-6.3788e-02],
            [ 1.8702e-01],
            [ 1.9750e-01],
            [-2.9756e-06],
            [-1.9489e-04],
            [ 1.2503e-03],
            [-2.1837e-04],
            [ 5.2441e-04],
            [-2.9308e-03],
            [-1.1519e-01],
            [-2.8608e-01],
            [-1.2138e-01],
            [-4.5510e-03],
            [-2.7456e-01],
            [-1.1996e-01],
            [ 4.5447e-02],
            [-6.6381e-02],
            [-1.5252e-02],
            [ 1.3836e-01],
            [ 1.7296e-01],
            [ 1.5306e-01],
            [-2.7365e-02],
            [-9.0870e-02],
            [-6.2779e-03],
            [-1.1479e-01],
            [-9.0850e-02],
            [-5.9928e-02],
            [-4.2558e-02],
            [ 9.1668e-02],
            [ 2.2334e-01],
            [ 1.9848e-01],
            [-1.1770e-04],
            [ 5.9704e-04],
            [-3.6128e-05],
            [-3.2918e-03],
            [-2.4367e-04],
            [-8.5253e-04],
            [-5.5297e-02],
            [-1.0907e-01],
            [ 9.3658e-02],
            [-3.9098e-02],
            [-3.5770e-02],
            [ 1.1277e-01],
            [ 3.8414e-02],
            [-5.3041e-02],
            [-1.7519e-01],
            [-6.2723e-02],
            [ 2.0702e-01],
            [ 2.3989e-01],
            [ 6.9063e-02],
            [ 2.9171e-01],
            [ 3.2693e-01],
            [ 1.8855e-01],
            [ 2.1947e-01],
            [ 3.9534e-02],
            [ 1.0615e-01],
            [ 2.2755e-01],
            [ 2.3191e-01],
            [ 1.3437e-01],
            [ 1.3848e-03],
            [ 8.0069e-04],
            [-1.1399e-03],
            [ 8.5605e-04],
            [-1.3270e-03],
            [ 3.0711e-03],
            [ 3.8801e-02],
            [ 1.5101e-01],
            [ 2.9244e-01],
            [ 2.6511e-01],
            [ 2.2122e-01],
            [ 1.8414e-01],
            [-1.2555e-01],
            [-2.3436e-01],
            [-3.1491e-01],
            [-5.5455e-02],
            [ 1.4551e-01],
            [ 2.8224e-02],
            [ 1.5079e-01],
            [ 2.9295e-01],
            [ 3.1877e-01],
            [ 2.4997e-01],
            [ 4.2725e-01],
            [ 2.9920e-01],
            [ 1.3090e-01],
            [ 1.3666e-01],
            [ 1.9278e-01],
            [ 4.7368e-02],
            [ 1.1104e-03],
            [ 3.7813e-04],
            [ 2.9895e-04],
            [ 1.1401e-03],
            [-3.4484e-05],
            [ 1.9998e-03],
            [ 1.4105e-01],
            [ 2.2985e-01],
            [ 1.3497e-01],
            [ 9.4158e-02],
            [ 1.6623e-01],
            [ 2.4428e-01],
            [-1.5544e-01],
            [-3.5388e-01],
            [-4.4814e-01],
            [-2.1033e-02],
            [-9.1014e-02],
            [-1.5735e-01],
            [-5.0167e-02],
            [ 7.4330e-02],
            [ 1.2126e-01],
            [ 9.6431e-02],
            [ 3.0829e-01],
            [ 2.0956e-01],
            [ 2.4590e-02],
            [ 1.6628e-02],
            [ 2.5059e-02],
            [ 8.9687e-03],
            [-6.3774e-04],
            [-4.6600e-04],
            [ 1.8111e-04],
            [ 9.1916e-04],
            [-8.9695e-04],
            [ 6.3894e-04],
            [ 4.7475e-02],
            [ 4.8938e-02],
            [-1.1028e-02],
            [-1.2694e-01],
            [-6.3835e-02],
            [ 2.3893e-01],
            [-7.2644e-02],
            [-2.7361e-01],
            [-1.9883e-01],
            [-6.1579e-02],
            [-2.0232e-01],
            [-7.0742e-02],
            [ 2.9617e-04],
            [ 2.2118e-02],
            [ 3.4191e-02],
            [ 6.3618e-02],
            [ 2.0117e-01],
            [ 3.2184e-02],
            [ 7.9059e-03],
            [-2.7789e-03],
            [ 1.3305e-04],
            [ 6.1062e-04],
            [-4.3401e-04],
            [-1.6789e-03],
            [-1.4672e-03],
            [ 4.7558e-04],
            [-1.6494e-04],
            [-1.1698e-04],
            [ 2.7396e-03],
            [-1.1052e-02],
            [-8.1558e-03],
            [-1.1144e-03],
            [-1.0480e-02],
            [ 8.3354e-02],
            [ 1.9237e-01],
            [ 4.9229e-02],
            [-5.0487e-02],
            [-2.3740e-02],
            [-1.5987e-01],
            [ 3.1901e-02],
            [ 6.0083e-02],
            [ 8.5258e-02],
            [ 1.9303e-02],
            [ 3.5519e-02],
            [ 7.8349e-03],
            [ 2.1035e-02],
            [ 1.0331e-02],
            [ 7.2469e-04],
            [-1.0284e-03],
            [ 1.2919e-04],
            [-7.2020e-04],
            [-9.1021e-04],
            [ 1.0591e-04],
            [ 1.6747e-03],
            [ 7.1057e-04],
            [ 8.8092e-04],
            [-4.4901e-04],
            [-8.1657e-04],
            [ 9.5063e-04],
            [ 5.1107e-03],
            [ 1.0821e-02],
            [ 2.5191e-03],
            [ 1.6329e-03],
            [ 7.2319e-04],
            [ 1.9945e-01],
            [ 1.6037e-01],
            [ 6.0999e-02],
            [ 8.4084e-02],
            [ 1.7198e-02],
            [-6.7118e-04],
            [-1.8027e-03],
            [ 1.1312e-03],
            [-1.3307e-03],
            [-3.5015e-04],
            [ 1.6007e-04],
            [ 2.2920e-03],
            [ 8.2785e-04],
            [ 2.0433e-04],
            [ 3.7831e-04],
            [-1.3899e-03],
            [-4.5663e-04],
            [ 2.0394e-03],
            [-5.7494e-04],
            [-1.7051e-03],
            [ 3.4849e-04],
            [ 2.8651e-04],
            [ 9.6406e-05],
            [-2.1055e-03],
            [ 7.2306e-04],
            [-1.5693e-03],
            [-7.3593e-04],
            [-7.9554e-05],
            [ 1.8915e-03],
            [-2.2229e-04],
            [ 1.4858e-02],
            [ 3.9102e-02],
            [ 1.3620e-02],
            [-1.4836e-04],
            [ 7.4546e-04],
            [-7.4074e-04],
            [-2.8766e-03],
            [ 1.2539e-04],
            [-1.3548e-04],
            [-1.0611e-03],
            [ 4.0102e-04],
            [-9.1278e-04],
            [-6.4459e-04],
            [ 3.9143e-05],
            [-6.4820e-04],
            [ 5.6154e-04],
            [-6.5917e-05],
            [-5.5342e-04],
            [ 8.5684e-04],
            [ 3.0157e-04],
            [-1.2654e-03],
            [-9.7653e-04],
            [-1.1835e-03],
            [-2.6084e-04],
            [-9.0921e-04],
            [-5.9687e-04],
            [ 7.5591e-04],
            [ 9.3186e-04],
            [-5.0229e-04],
            [-6.3767e-04],
            [-9.6996e-04],
            [ 1.9786e-03],
            [ 1.7232e-03],
            [ 4.5464e-04],
            [-9.9913e-04],
            [-7.4080e-04],
            [ 3.1478e-04],
            [ 9.6729e-04],
            [-1.8139e-03],
            [-1.5460e-04],
            [-3.5329e-04],
            [ 2.6354e-04]])
    [stdout:3]
    [DG] Rank 3: Global Gradient Norm: 0.00021530537924263626
    tensor([[ 1.1936e-03],
            [-7.4490e-04],
            [-1.8730e-03],
            [-3.7072e-04],
            [ 2.8816e-04],
            [ 1.0057e-03],
            [ 2.2110e-04],
            [-2.9943e-04],
            [-4.1486e-04],
            [-6.4886e-04],
            [ 4.9854e-05],
            [-7.5339e-04],
            [ 5.8118e-04],
            [ 9.4398e-04],
            [ 6.8303e-05],
            [-6.3700e-04],
            [ 4.4353e-04],
            [-1.3363e-04],
            [-1.7060e-04],
            [-1.5143e-03],
            [ 2.2467e-05],
            [-1.9249e-04],
            [-5.2772e-04],
            [ 9.2492e-04],
            [ 1.2930e-03],
            [ 3.8942e-04],
            [-1.9009e-04],
            [ 3.2923e-04],
            [-7.6191e-04],
            [ 7.4583e-04],
            [-2.3638e-03],
            [ 3.3711e-04],
            [-1.1406e-03],
            [-3.9181e-04],
            [ 1.5934e-03],
            [ 8.3065e-04],
            [ 1.9899e-04],
            [ 2.1024e-04],
            [-1.3065e-03],
            [-1.6548e-03],
            [ 9.3925e-04],
            [ 1.4836e-04],
            [ 1.8946e-03],
            [-3.9512e-03],
            [-1.4238e-02],
            [-5.6233e-03],
            [ 2.1988e-03],
            [-5.7531e-05],
            [ 3.8517e-04],
            [ 2.2053e-04],
            [-3.1304e-04],
            [ 7.5114e-04],
            [-9.8908e-05],
            [-1.4486e-04],
            [ 8.3906e-05],
            [-2.6352e-04],
            [-2.1619e-04],
            [-4.1876e-04],
            [-3.8936e-04],
            [ 1.3845e-03],
            [ 9.1240e-04],
            [-7.6632e-04],
            [-7.6944e-04],
            [-4.9053e-04],
            [-4.7654e-04],
            [-2.2263e-04],
            [ 1.0965e-03],
            [ 9.7702e-04],
            [ 8.0908e-03],
            [ 1.0465e-01],
            [ 3.1371e-02],
            [ 1.7986e-02],
            [ 6.2185e-02],
            [ 1.3780e-01],
            [ 3.5558e-02],
            [ 2.8015e-03],
            [ 9.3000e-04],
            [-4.3380e-04],
            [ 3.2760e-04],
            [ 7.6509e-04],
            [-9.6313e-04],
            [-1.2922e-03],
            [-8.9881e-04],
            [ 1.0799e-03],
            [ 9.1239e-04],
            [-1.5100e-04],
            [-4.6868e-04],
            [-2.2949e-05],
            [-9.2020e-04],
            [-6.8143e-04],
            [-2.0873e-04],
            [ 8.0167e-04],
            [ 1.0338e-02],
            [ 1.1304e-02],
            [ 2.6729e-03],
            [ 1.4059e-03],
            [ 5.7374e-02],
            [ 1.1832e-01],
            [ 1.1785e-01],
            [ 1.2431e-01],
            [ 2.3523e-01],
            [ 2.6790e-01],
            [ 1.2935e-01],
            [ 1.8793e-02],
            [ 7.3079e-03],
            [ 3.3638e-03],
            [ 1.9628e-04],
            [-8.9218e-05],
            [ 1.6256e-03],
            [-8.4801e-04],
            [-8.4250e-04],
            [-2.0151e-03],
            [ 1.1065e-03],
            [-3.5768e-04],
            [-1.1408e-03],
            [ 3.1309e-04],
            [-4.0829e-04],
            [ 9.1734e-04],
            [-1.8682e-03],
            [-4.7004e-03],
            [ 8.4736e-02],
            [ 1.3099e-01],
            [ 4.1886e-02],
            [ 4.9341e-02],
            [ 1.1376e-01],
            [ 1.2543e-01],
            [ 5.5223e-02],
            [ 1.2273e-01],
            [ 7.4710e-02],
            [ 1.9841e-01],
            [-9.3371e-02],
            [ 3.6030e-02],
            [ 1.5696e-01],
            [ 1.0093e-01],
            [ 7.7009e-02],
            [ 3.7515e-02],
            [ 6.0213e-03],
            [ 2.0485e-03],
            [-2.5639e-04],
            [-9.6560e-04],
            [ 1.7806e-04],
            [ 1.0697e-03],
            [ 2.9464e-04],
            [ 1.8435e-03],
            [-2.8556e-05],
            [ 5.0634e-03],
            [-3.1504e-03],
            [-5.5291e-03],
            [ 1.2589e-01],
            [ 1.4413e-01],
            [ 3.3708e-02],
            [ 5.0546e-02],
            [ 1.7358e-01],
            [ 2.8364e-01],
            [-4.6394e-02],
            [-1.4584e-01],
            [ 7.1890e-04],
            [ 6.6287e-02],
            [-1.4319e-01],
            [-5.7232e-02],
            [ 2.0920e-01],
            [ 1.8760e-01],
            [ 1.4816e-01],
            [ 1.3011e-01],
            [ 4.9874e-02],
            [ 2.9479e-04],
            [-1.2428e-03],
            [ 7.0202e-04],
            [ 1.7221e-04],
            [-4.5142e-04],
            [-1.2387e-03],
            [-3.8823e-04],
            [ 3.2491e-04],
            [ 6.6376e-04],
            [ 1.7596e-02],
            [ 7.2569e-02],
            [ 8.0967e-02],
            [ 2.5680e-02],
            [-3.8500e-02],
            [-6.2019e-02],
            [ 7.6648e-02],
            [ 1.1012e-01],
            [-5.1055e-02],
            [-1.1972e-01],
            [-1.4509e-01],
            [-6.6121e-02],
            [-5.4169e-02],
            [-1.0989e-01],
            [ 2.1135e-01],
            [ 2.1979e-01],
            [ 1.0230e-02],
            [ 3.2104e-02],
            [ 4.3940e-02],
            [-2.8970e-03],
            [-1.8488e-03],
            [-3.3765e-04],
            [ 1.5543e-03],
            [ 1.3449e-04],
            [ 4.7694e-04],
            [-7.3921e-05],
            [ 1.3149e-01],
            [ 1.4051e-01],
            [ 6.2532e-02],
            [ 1.4592e-01],
            [ 1.6816e-02],
            [-1.5681e-01],
            [-2.4907e-01],
            [-2.8624e-01],
            [-4.3275e-02],
            [ 4.6067e-02],
            [-7.0588e-03],
            [-2.9414e-02],
            [ 1.1140e-02],
            [-1.7544e-01],
            [-1.3208e-01],
            [-1.4081e-01],
            [ 5.3895e-02],
            [-5.3165e-02],
            [-1.0676e-01],
            [-5.8235e-02],
            [-6.4296e-02],
            [-9.1896e-03],
            [ 1.4059e-04],
            [-1.9690e-04],
            [-4.9813e-05],
            [ 2.7041e-04],
            [ 4.3805e-04],
            [ 3.2821e-04],
            [ 1.2157e-01],
            [ 1.6255e-01],
            [ 5.8713e-02],
            [ 1.1099e-01],
            [ 1.1229e-02],
            [-1.6565e-01],
            [-2.0318e-01],
            [-1.4503e-01],
            [ 1.4282e-01],
            [ 3.0895e-03],
            [ 2.4932e-01],
            [ 1.5670e-01],
            [-2.0866e-01],
            [-2.5122e-01],
            [ 1.4685e-03],
            [-1.7111e-02],
            [-1.9777e-01],
            [-2.7216e-01],
            [-1.5315e-01],
            [-6.2656e-02],
            [-3.4171e-02],
            [-4.0108e-03],
            [ 2.4435e-03],
            [-5.9737e-04],
            [ 3.4371e-04],
            [-1.5111e-03],
            [ 1.0139e-04],
            [ 1.4529e-03],
            [ 1.4834e-01],
            [ 1.3642e-01],
            [-4.3561e-02],
            [ 4.7752e-02],
            [ 1.8911e-02],
            [-3.8652e-02],
            [-8.7196e-02],
            [-7.6051e-03],
            [ 1.2670e-01],
            [ 2.7121e-01],
            [ 3.0166e-01],
            [ 1.1256e-01],
            [-1.5079e-02],
            [-9.3460e-02],
            [ 1.6718e-01],
            [ 1.2887e-01],
            [-2.8524e-01],
            [-3.9803e-01],
            [-2.4504e-01],
            [-1.3680e-01],
            [-2.6084e-02],
            [-5.9157e-03],
            [ 2.1796e-05],
            [-1.4788e-05],
            [ 1.6040e-03],
            [ 2.4874e-04],
            [-1.2744e-03],
            [-2.0439e-04],
            [ 1.4752e-01],
            [ 3.9524e-02],
            [-1.2499e-01],
            [-1.2120e-01],
            [-1.1146e-01],
            [-1.6405e-02],
            [-3.6145e-02],
            [ 4.2552e-02],
            [ 1.1004e-01],
            [ 1.1221e-01],
            [ 1.1547e-01],
            [ 2.6752e-01],
            [ 2.1524e-01],
            [-1.0571e-01],
            [ 1.5370e-02],
            [-2.1486e-02],
            [-3.0315e-01],
            [-3.4109e-01],
            [-2.8405e-01],
            [-1.5763e-01],
            [-4.4092e-02],
            [-3.7633e-03],
            [-3.5337e-04],
            [ 8.8794e-04],
            [ 3.6015e-04],
            [-2.2263e-03],
            [ 6.5641e-04],
            [-3.8196e-04],
            [ 1.3840e-01],
            [-3.8922e-02],
            [-1.6563e-01],
            [-1.9044e-01],
            [-2.2366e-01],
            [-1.6649e-01],
            [-1.2419e-01],
            [-9.3218e-02],
            [-7.1844e-03],
            [ 4.7069e-02],
            [ 1.2733e-01],
            [ 5.0493e-01],
            [ 3.3510e-01],
            [-2.4408e-02],
            [-1.7265e-01],
            [-2.5690e-01],
            [-1.1942e-01],
            [-2.5489e-01],
            [-2.5832e-01],
            [-9.0952e-02],
            [-7.7847e-02],
            [-7.0345e-03],
            [ 3.9934e-05],
            [-1.2058e-03],
            [-4.2598e-04],
            [-1.5268e-03],
            [ 1.6235e-03],
            [-3.5158e-05],
            [ 7.6744e-02],
            [-1.2663e-01],
            [-1.7769e-01],
            [-1.6340e-01],
            [-2.7658e-01],
            [-2.9604e-01],
            [-1.2244e-01],
            [-1.3711e-01],
            [-2.3412e-01],
            [ 1.6656e-01],
            [ 3.7587e-01],
            [ 6.8338e-01],
            [ 3.1691e-01],
            [-8.0374e-03],
            [-3.2811e-01],
            [-3.2229e-01],
            [-7.1799e-02],
            [-2.9149e-01],
            [-2.6707e-01],
            [-1.7684e-01],
            [-9.4497e-02],
            [-1.0763e-02],
            [ 2.3799e-03],
            [ 3.4798e-04],
            [ 1.6462e-04],
            [ 1.6323e-04],
            [ 3.9174e-04],
            [-6.8513e-04],
            [ 1.7739e-03],
            [-1.6999e-01],
            [-1.7110e-01],
            [-2.7079e-01],
            [-3.4341e-01],
            [-2.2118e-01],
            [-1.7788e-01],
            [-3.4913e-01],
            [-1.6539e-01],
            [ 4.6551e-01],
            [ 7.1713e-01],
            [ 6.3820e-01],
            [ 2.7224e-01],
            [-7.4490e-02],
            [-6.6655e-01],
            [-2.4464e-01],
            [ 2.1602e-03],
            [-4.0832e-01],
            [-4.2312e-01],
            [-2.6096e-01],
            [-1.0566e-01],
            [-1.3835e-02],
            [-1.9651e-04],
            [-1.4117e-03],
            [ 4.5764e-05],
            [ 5.3655e-05],
            [-1.3077e-04],
            [-7.7662e-04],
            [-3.7123e-02],
            [-1.8083e-01],
            [-2.5239e-01],
            [-3.6627e-01],
            [-2.5280e-01],
            [-6.4903e-02],
            [-4.1040e-01],
            [-3.1260e-01],
            [-2.3981e-02],
            [ 5.3778e-01],
            [ 8.9399e-01],
            [ 5.5042e-01],
            [ 2.8988e-01],
            [-2.5772e-01],
            [-8.2325e-01],
            [-1.4288e-01],
            [-1.4194e-01],
            [-4.0810e-01],
            [-3.4428e-01],
            [-2.6352e-01],
            [-1.1779e-01],
            [-1.5645e-02],
            [-1.0376e-03],
            [ 3.3638e-04],
            [ 5.3145e-04],
            [ 9.4312e-05],
            [ 1.2118e-03],
            [ 2.3385e-03],
            [-1.0867e-01],
            [-2.2869e-01],
            [-2.7879e-01],
            [-2.9921e-01],
            [-2.1537e-01],
            [-1.9712e-01],
            [-4.9786e-01],
            [-2.4333e-01],
            [ 2.0434e-01],
            [ 7.5539e-01],
            [ 9.3154e-01],
            [ 3.3754e-01],
            [ 3.2506e-02],
            [-5.7489e-01],
            [-4.9128e-01],
            [-1.3184e-01],
            [-2.0079e-01],
            [-3.4799e-01],
            [-2.8815e-01],
            [-2.6709e-01],
            [-1.1316e-01],
            [-1.4170e-02],
            [ 1.2462e-04],
            [ 1.3099e-03],
            [ 1.1588e-03],
            [-8.7237e-04],
            [ 7.3758e-04],
            [ 7.3033e-04],
            [-1.2493e-01],
            [-2.2329e-01],
            [-2.4406e-01],
            [-3.1968e-01],
            [-1.8771e-01],
            [-4.9374e-01],
            [-4.0576e-01],
            [-6.5834e-02],
            [ 2.2978e-01],
            [ 7.2183e-01],
            [ 6.0289e-01],
            [ 2.8153e-01],
            [-2.3492e-01],
            [-5.7714e-01],
            [-1.5804e-01],
            [-1.4676e-01],
            [-3.6442e-01],
            [-2.7534e-01],
            [-2.5881e-01],
            [-2.0027e-01],
            [ 4.8282e-02],
            [ 3.9790e-02],
            [-4.6834e-04],
            [ 8.8627e-04],
            [-5.8268e-04],
            [ 5.1061e-04],
            [-1.0266e-05],
            [-5.9888e-04],
            [-1.3734e-01],
            [-2.4802e-01],
            [-2.4900e-01],
            [-2.3358e-01],
            [-2.5944e-01],
            [-5.1954e-01],
            [-3.3076e-01],
            [-3.4832e-02],
            [ 2.4520e-01],
            [ 4.9237e-01],
            [ 2.3397e-01],
            [ 2.2124e-01],
            [-2.1235e-01],
            [-4.4434e-01],
            [-2.2293e-01],
            [-2.3142e-01],
            [-3.1346e-01],
            [-1.9325e-01],
            [-2.2575e-01],
            [-1.5276e-01],
            [ 1.2380e-01],
            [ 1.9340e-01],
            [ 1.6610e-04],
            [-6.1020e-04],
            [ 3.0097e-04],
            [ 1.4370e-03],
            [-1.0996e-03],
            [-2.1737e-03],
            [-1.2868e-01],
            [-2.8049e-01],
            [-2.2127e-01],
            [-1.2364e-01],
            [-3.0615e-01],
            [-3.7004e-01],
            [-1.5353e-01],
            [-1.4771e-01],
            [ 1.8286e-01],
            [ 2.9109e-01],
            [ 1.3559e-01],
            [ 1.0926e-01],
            [-3.2381e-02],
            [-3.4588e-01],
            [-8.7224e-02],
            [-2.9529e-01],
            [-1.4865e-01],
            [-1.4098e-01],
            [-1.6710e-01],
            [-6.4363e-02],
            [ 1.8519e-01],
            [ 1.9853e-01],
            [ 9.9199e-04],
            [ 1.1208e-03],
            [ 4.5193e-04],
            [ 1.1877e-04],
            [-2.9347e-03],
            [-1.4341e-03],
            [-1.1719e-01],
            [-2.8667e-01],
            [-1.1880e-01],
            [-4.7903e-03],
            [-2.7615e-01],
            [-1.2099e-01],
            [ 4.5040e-02],
            [-6.6462e-02],
            [-1.5354e-02],
            [ 1.4131e-01],
            [ 1.7422e-01],
            [ 1.5415e-01],
            [-2.7758e-02],
            [-8.9533e-02],
            [-5.3437e-03],
            [-1.1659e-01],
            [-8.9402e-02],
            [-6.0854e-02],
            [-4.1805e-02],
            [ 8.8924e-02],
            [ 2.2384e-01],
            [ 1.9995e-01],
            [-1.0703e-04],
            [ 1.7625e-04],
            [ 4.8137e-04],
            [-1.4263e-03],
            [ 7.7820e-04],
            [-2.1040e-03],
            [-5.5054e-02],
            [-1.1095e-01],
            [ 9.4814e-02],
            [-4.1036e-02],
            [-3.4830e-02],
            [ 1.1223e-01],
            [ 3.9105e-02],
            [-5.2370e-02],
            [-1.7553e-01],
            [-6.0741e-02],
            [ 2.1068e-01],
            [ 2.4118e-01],
            [ 6.8158e-02],
            [ 2.9158e-01],
            [ 3.2611e-01],
            [ 1.8905e-01],
            [ 2.1936e-01],
            [ 3.9863e-02],
            [ 1.0905e-01],
            [ 2.2599e-01],
            [ 2.3342e-01],
            [ 1.3539e-01],
            [ 4.6248e-04],
            [ 3.7010e-05],
            [-7.4634e-05],
            [ 1.1048e-03],
            [-6.3754e-04],
            [ 3.9671e-03],
            [ 3.9264e-02],
            [ 1.5106e-01],
            [ 2.9059e-01],
            [ 2.6630e-01],
            [ 2.2392e-01],
            [ 1.8512e-01],
            [-1.2532e-01],
            [-2.3419e-01],
            [-3.1322e-01],
            [-5.6272e-02],
            [ 1.4562e-01],
            [ 2.8255e-02],
            [ 1.4950e-01],
            [ 2.9165e-01],
            [ 3.1740e-01],
            [ 2.5070e-01],
            [ 4.2637e-01],
            [ 2.9747e-01],
            [ 1.3318e-01],
            [ 1.3692e-01],
            [ 1.9404e-01],
            [ 4.6632e-02],
            [-5.7465e-04],
            [ 1.9558e-03],
            [ 9.9975e-04],
            [-1.2616e-03],
            [ 3.9367e-05],
            [ 1.2830e-03],
            [ 1.3852e-01],
            [ 2.3072e-01],
            [ 1.3232e-01],
            [ 9.5084e-02],
            [ 1.6703e-01],
            [ 2.4283e-01],
            [-1.5539e-01],
            [-3.5331e-01],
            [-4.4638e-01],
            [-2.2979e-02],
            [-9.1276e-02],
            [-1.5958e-01],
            [-4.9901e-02],
            [ 7.3891e-02],
            [ 1.2351e-01],
            [ 9.5988e-02],
            [ 3.0846e-01],
            [ 2.0827e-01],
            [ 2.6555e-02],
            [ 1.5219e-02],
            [ 2.5040e-02],
            [ 1.0436e-02],
            [ 1.7145e-03],
            [-7.1479e-04],
            [-1.2700e-04],
            [ 5.8395e-04],
            [ 1.6302e-03],
            [-5.7736e-04],
            [ 4.8043e-02],
            [ 4.7976e-02],
            [-1.2613e-02],
            [-1.2770e-01],
            [-6.7093e-02],
            [ 2.3959e-01],
            [-7.2696e-02],
            [-2.7296e-01],
            [-2.0017e-01],
            [-6.1485e-02],
            [-2.0350e-01],
            [-7.3806e-02],
            [ 7.1825e-04],
            [ 2.2464e-02],
            [ 3.4759e-02],
            [ 6.1107e-02],
            [ 2.0083e-01],
            [ 3.1485e-02],
            [ 5.8003e-03],
            [ 1.5384e-03],
            [-4.9454e-04],
            [-4.6774e-04],
            [-1.8948e-03],
            [ 1.2472e-03],
            [-6.4507e-04],
            [-4.6619e-04],
            [-9.1496e-05],
            [ 4.6719e-04],
            [-1.1312e-03],
            [-1.0250e-02],
            [-6.3359e-03],
            [-1.0306e-03],
            [-1.0000e-02],
            [ 8.2875e-02],
            [ 1.9323e-01],
            [ 5.0185e-02],
            [-5.0764e-02],
            [-2.5343e-02],
            [-1.6012e-01],
            [ 3.2788e-02],
            [ 6.1196e-02],
            [ 8.5378e-02],
            [ 2.1177e-02],
            [ 3.6249e-02],
            [ 6.1352e-03],
            [ 2.0501e-02],
            [ 1.1743e-02],
            [ 8.5186e-04],
            [-1.7395e-04],
            [-1.3267e-03],
            [-1.3114e-03],
            [-4.1113e-04],
            [ 2.6471e-04],
            [ 7.8780e-04],
            [-1.1546e-03],
            [ 9.8577e-04],
            [-1.6558e-04],
            [ 5.7647e-04],
            [ 1.2162e-04],
            [ 4.1080e-03],
            [ 1.0549e-02],
            [ 4.4950e-03],
            [-2.3869e-04],
            [ 2.7160e-03],
            [ 2.0082e-01],
            [ 1.6118e-01],
            [ 6.1469e-02],
            [ 8.4072e-02],
            [ 1.7520e-02],
            [-3.5401e-04],
            [ 5.7929e-04],
            [-6.7392e-04],
            [-8.9712e-04],
            [-6.6057e-04],
            [-1.5794e-03],
            [ 1.5097e-03],
            [ 2.7185e-04],
            [-5.7224e-05],
            [ 1.1565e-03],
            [ 9.9847e-05],
            [-4.0673e-04],
            [-8.9043e-05],
            [ 1.5420e-03],
            [-1.0391e-03],
            [ 7.4576e-04],
            [-1.6643e-03],
            [ 1.9407e-03],
            [ 1.2712e-03],
            [-1.1068e-03],
            [ 5.3735e-04],
            [ 1.9774e-03],
            [ 6.7940e-04],
            [ 1.2226e-04],
            [ 1.3491e-05],
            [ 1.3944e-02],
            [ 3.8272e-02],
            [ 1.1215e-02],
            [-5.0589e-04],
            [-9.3688e-04],
            [ 1.2592e-03],
            [-1.7366e-04],
            [ 1.0892e-03],
            [ 3.5814e-04],
            [ 3.2523e-04],
            [-1.2858e-03],
            [ 6.3742e-04],
            [ 1.4546e-04],
            [-4.5635e-05],
            [-5.4926e-04],
            [ 7.8729e-06],
            [-1.0660e-03],
            [-5.6333e-04],
            [-1.3632e-03],
            [ 5.9416e-04],
            [ 1.2009e-04],
            [ 6.0085e-05],
            [-4.1247e-05],
            [ 1.5440e-03],
            [-2.6223e-04],
            [-1.5850e-03],
            [-1.4930e-03],
            [ 4.0325e-05],
            [-1.7284e-03],
            [ 8.2065e-04],
            [ 1.1659e-03],
            [-3.0242e-04],
            [-1.3712e-03],
            [-1.9559e-03],
            [-2.7808e-04],
            [-8.4462e-04],
            [ 1.5066e-03],
            [-4.1543e-04],
            [ 1.5120e-03],
            [-4.9619e-04],
            [-2.5322e-04],
            [ 1.1414e-04]])

```
# Generate topology.

# Plot figure
%px G = topology_util.ExponentialTwoGraph(bf.size())

G_0 = dview.pull("G", block=True, targets=0)
nx.draw_circular(G_0)
```

Run the Adapt-Then-Combine Algorithm.

```
%%px

# Set topology as exponential-2 graph.
bf.set_topology(G)
topology = bf.load_topology()
self_weight, neighbor_weights = topology_util.GetRecvWeights(topology, bf.rank())
```

```
%%px
# Print weights.

if bf.rank() == 0:
    print("self weights: {}\n".format(self_weight))
    print("neighbor weights:")
    for k, v in neighbor_weights.items():
        print(k, v)
```

    [stdout:0]
    self weights: 0.3333333333333333

    neighbor weights:
    2 0.3333333333333333
    3 0.3333333333333333

```
# Run the Adapt-Then-Combine Algorithm

maxite = 1000
relative_error = torch.zeros((maxite, 1))
%px w = None
for ite in range(maxite):

    if ite%100 == 0:
        print('Progress {}/{}'.format(ite, maxite))

    %px w, relative_error_ite = ATC(subset_data, subset_targets, w_opt_mnist, step_size=1, w=w)
    relative_error[ite] = dview.pull('relative_error_ite', block=True, targets=0)
```

    Progress 0/1000
    Progress 100/1000
    Progress 200/1000
    Progress 300/1000
    Progress 400/1000
    Progress 500/1000
    Progress 600/1000
    Progress 700/1000
    Progress 800/1000
    Progress 900/1000

```
plt.semilogy(relative_error)
```

```
# Run the Adapt-With-Combine Algorithm

maxite = 1000
relative_error = torch.zeros((maxite, 1))
%px w = None
for ite in range(maxite):

    if ite%100 == 0:
        print('Progress {}/{}'.format(ite, maxite))

    %px w, relative_error_ite = AWC(subset_data, subset_targets, w_opt_mnist, step_size=1, w=w)
    relative_error[ite] = dview.pull('relative_error_ite', block=True, targets=0)
```

    Progress 0/1000
    Progress 100/1000
    Progress 200/1000
    Progress 300/1000
    Progress 400/1000
    Progress 500/1000
    Progress 600/1000
    Progress 700/1000
    Progress 800/1000
    Progress 900/1000

```
plt.semilogy(relative_error)
```

For our plots, we used relative error which is the L2-norm of the difference between the current set of parameters and the optimal set of parameters divided by the L2-norm of the optimal set of parameters.