pytorch使用tpu

!pip install torch_xla

import torch
import torch_xla.core.xla_model as xm
import torch.nn as nn
import torch.optim as optim

device = xm.xla_device()

model = model.to(device)

criterion = nn.CrossEntropyLoss().to(device)
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9).to(device)

for epoch in range(num_epochs):
    model.train()
    for i, (inputs, labels) in enumerate(train_loader):
        inputs = inputs.to(device)
        labels = labels.to(device)
        # 正向传播
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    model.eval()
    with torch.no_grad():
        # 推理代码