File size: 503 Bytes
0ba62b0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a791472
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
<pre>
+from accelerate import Accelerator
+accelerator = Accelerator()
+dataloader, model, optimizer scheduler = accelerator.prepare(
+        dataloader, model, optimizer, scheduler
+)

for batch in dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
-    inputs = inputs.to(device)
-    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
-    loss.backward()
+    accelerator.backward(loss)
    optimizer.step()
    scheduler.step()</pre>