import tensorflow as tf from tensorflow.keras import optimizers, layers, Model class linearModel(Model): def __init__(self, ndim): super(linearModel, self).__init__() self.w = tf.Variable( shape=[ndim, 1], initial_value=tf.random.uniform( [ndim,1], minval=-0.1, maxval=0.1, dtype=tf.float32)) @tf.function def call(self, x): y = tf.squeeze(tf.matmul(x, self.w), axis=1) return y (xs, ys), (o_x, o_y) = load_data('train.txt') ndim = xs.shape[1] model = linearModel(ndim=ndim) 将上述代码的tensflow框架改为pytorch框架，并能在jupyter上运行

import torch
import torch.nn as nn
import torch.optim as optim

class LinearModel(nn.Module):
def __init__(self, ndim):
super(LinearModel, self).__init__()
self.w = nn.Parameter(torch.rand(ndim, 1) * 0.2 - 0.1)

def forward(self, x):
y = torch.mm(x, self.w).squeeze(dim=1)
return y

def load_data(file_path):
# 加载数据的代码请自行实现
pass

# 加载数据
(xs, ys), (o_x, o_y) = load_data('train.txt')
ndim = xs.shape[1]

# 定义模型、损失函数和优化器
model = LinearModel(ndim=ndim)
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练模型
for epoch in range(1000):
optimizer.zero_grad()
y_pred = model(xs)
loss = criterion(y_pred, ys)
loss.backward()
optimizer.step()

if epoch % 100 == 0:
print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, 1000, loss.item()))

# 测试模型
with torch.no_grad():
y_test = model(o_x)
test_loss = criterion(y_test, o_y)
print('Test Loss: {:.4f}'.format(test_loss.item()))