我想测试下gpu进行ai运算可以帮我找一个简单的代码吗

当然可以。在Python中，特别是在深度学习框架如PyTorch或TensorFlow中，你可以使用GPU来进行AI运算。这里是一个简单的例子，我们将使用PyTorch来展示如何加载数据到GPU并执行计算：

# 首先安装必要的库（如果还没安装）
!pip install torch

import torch
import torch.nn as nn
from torch.autograd import Variable

# 检查设备是否可用（默认使用GPU，如果没设置则会用CPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# 定义一个小神经网络模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 1)  # 假设输入10维，输出1维

    def forward(self, x):
        return self.fc(x)

model = SimpleModel().to(device)  # 将模型移到指定设备

# 假设我们有样本数据
x_data = torch.randn(100, 10).to(device)  # 创建随机数据 (batch_size, input_dim)
y_data = torch.randn(100, 1).to(device)   # 目标标签

# 将数据视为Variable以便自动求导
inputs = Variable(x_data)
targets = Variable(y_data)

# 执行一次前向传播
outputs = model(inputs)
loss = nn.MSELoss()(outputs, targets)  # 使用均方误差作为损失函数

# GPU运算示例 - 计算梯度
loss.backward()

print("Simple AI computation on GPU completed.")