I forgot to set the rank for GPU here in this code: import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
import torch.distributed as dist
import itertools

# Initialize Distributed Training
dist.init_process_group(backend=“nccl”)
device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)
torch.cuda.set_device(dist.get_rank())

# Define Training Hyperparameters
batch_size = 128
learning_rate = 0.001
epochs = 10
patience = None
scheduler_patience = None
optimization_technique = ‘early_stopping’

num_workers = 4
prefetch_factor = 2
persistent_workers = True

accumulation_steps = 4

# Define Model
model = …

# Define Data Loaders
image_dataset = …
train_sampler = DistributedSampler(image_dataset, num_replicas=dist.get_world_size(), rank=dist.get_rank())
train_loader = DataLoader(
image_dataset,
batch_size=batch_size,
sampler=train_sampler,
num_workers=num_workers,
pin_memory=True,
prefetch_factor=prefetch_factor,
persistent_workers=persistent_workers
)

test_loader = DataLoader(
image_dataset,
batch_size=batch_size,
sampler=torch.utils.data.SubsetRandomSampler(test_indices),
num_workers=num_workers,
pin_memory=True,
prefetch_factor=prefetch_factor,
persistent_workers=persistent_workers
)


# Define Training Loop
model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

if optimization_technique == ‘learning_rate_scheduler’ and scheduler_patience:
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, ‘min’, patience=scheduler_patience, verbose=True)
else:
scheduler = None

train_losses = []
train_accuracies = []

test_losses = []
test_accuracies = []

best_loss = float(‘inf’)
stopping_counter = 0

# Enable cudnn benchmark
torch.backends.cudnn.benchmark = True

# Enable Automatic Mixed Precision (AMP)
scaler = torch.cuda.amp.GradScaler()

for epoch in range(epochs):
train_stream = torch.cuda.Stream()
test_stream = torch.cuda.Stream()

running_train_loss = 0.0
running_train_acc = 0
num_batches_train = 0

optimizer.zero_grad(set_to_none=True)

for i, (X_batch, y_batch) in enumerate(train_loader):
# Move batch to device
X_batch, y_batch = X_batch.to(device, non_blocking=True), y_batch.to(device, non_blocking=True)

# Start asynchronous data transfer
with torch.cuda.stream(train_stream):
# Automatic Mixed Precision (AMP)
with torch.cuda.amp.autocast():
y_pred = model(X_batch.half())
loss = criterion(y_pred, y_batch)

scaler.scale(loss/accumulation_steps).backward()

if (i + 1) % accumulation_steps == 0:
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)

running_train_loss += loss.item()
running_train_acc += torch.sum(torch.eq(y_pred.argmax(dim=1), y_batch)).item()
num_batches_train += 1
train_stream.synchronize()

train_losses.append(running_train_loss / num_batches_train)
train_accuracies.append(running_train_acc / len(train_loader.dataset))

# Testing segment
running_test_loss = 0.0
running_test_acc = 0
num_batches_test = 0

with torch.no_grad():
for X_batch, y_batch in test_loader:

# Move batch to device
X_batch, y_batch = X_batch.to(device, non_blocking=True), y_batch.to(device, non_blocking=True)

# Start asynchronous data transfer
with torch.cuda.stream(test_stream):
with torch.cuda.amp.autocast():
y_pred = model(X_batch.half())
loss = criterion(y_pred, y_batch)

running_test_loss += loss.item()
running_test_acc += torch.sum(torch.eq(y_pred.argmax(dim=1), y_batch)).item()
num_batches_test += 1
test_stream.synchronize()

test_losses.append(running_test_loss / num_batches_test)
test_accuracies.append(running_test_acc / len(test_loader.dataset))

# Early stopping
if optimization_technique == ‘early_stopping’ and patience:
if test_losses[-1] < best_loss:
best_loss = test_losses[-1]
stopping_counter = 0
else:
stopping_counter += 1
if stopping_counter > patience:
print(f"Early stopping at epoch {epoch+1}/{epochs}")
break

# Learning rate scheduler
if optimization_technique == ‘learning_rate_scheduler’ and scheduler_patience and scheduler:
scheduler.step(test_losses[-1])

print(f"Epoch: {epoch+1}/{epochs}, Training Loss: {train_losses[-1]}, Test Loss: {test_losses[-1]}, Training Accuracy: {train_accuracies