import torch
import torch.nn as nn

accumulation_steps = 4   # 4번 누적 → 유효 배치 크기 4배
batch_size = 8           # 실제 미니배치 크기
# 유효 배치 크기 = 8 × 4 = 32

model = nn.Linear(784, 10).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

optimizer.zero_grad()    # 루프 밖에서 한 번 초기화

for step, (data, labels) in enumerate(train_loader):
    data, labels = data.to(device), labels.to(device)

    # 순전파 + 역전파 (기울기 누적)
    outputs = model(data)
    loss = criterion(outputs, labels)
    loss = loss / accumulation_steps  # 손실 스케일링
    loss.backward()

    # N 스텝마다 가중치 업데이트
    if (step + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

from torch.amp import autocast, GradScaler

scaler = GradScaler('cuda')
accumulation_steps = 4

optimizer.zero_grad()

for step, (data, labels) in enumerate(train_loader):
    data, labels = data.to(device), labels.to(device)

    with autocast('cuda'):
        outputs = model(data)
        loss = criterion(outputs, labels) / accumulation_steps

    scaler.scale(loss).backward()

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