基于pytorch，编写一段WGAN的代码

当然可以！以下是一个用PyTorch编写的简单WGAN代码示例：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import grad
from torch.utils.data import DataLoader

# Generator Model
class Generator(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(Generator, self).__init__()
        self.fc = nn.Linear(input_dim, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, output_dim)
    
    def forward(self, x):
        x = self.fc(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# Critic Model
class Critic(nn.Module):
    def __init__(self, input_dim):
        super(Critic, self).__init__()
        self.fc = nn.Linear(input_dim, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 1)
    
    def forward(self, x):
        x = self.fc(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# WGAN Trainer
class WGAN:
    def __init__(self, input_dim, output_dim, critic_iterations=5, batch_size=64, learning_rate=0.0001):
        self.generator = Generator(input_dim, output_dim)
        self.critic = Critic(output_dim)
        
        self.critic_iterations = critic_iterations
        self.batch_size = batch_size
        
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        self.generator.to(self.device)
        self.critic.to(self.device)
        
        self.generator_optimizer = optim.RMSprop(self.generator.parameters(), lr=learning_rate)
        self.critic_optimizer = optim.RMSprop(self.critic.parameters(), lr=learning_rate)
    
    def train(self, data_loader, num_epochs):
        for epoch in range(num_epochs):
            for i, real_data in enumerate(data_loader):
                real_data = real_data.to(self.device)
                
                # Critic Training
                for _ in range(self.critic_iterations):
                    self.critic_optimizer.zero_grad()
                    
                    noise = torch.randn(self.batch_size, input_dim).to(self.device)
                    fake_data = self.generator(noise)
                    
                    critic_loss = self.critic(fake_data).mean() - self.critic(real_data).mean()
                    critic_loss.backward()
                    
                    self.critic_optimizer.step()
                    
                    # Clip critic weights
                    for p in self.critic.parameters():
                        p.data.clamp_(-0.01, 0.01)
                
                # Generator Training
                self.generator_optimizer.zero_grad()
                
                noise = torch.randn(self.batch_size, input_dim).to(self.device)
                fake_data = self.generator(noise)
                
                generator_loss = -self.critic(fake_data).mean()
                generator_loss.backward()
                
                self.generator_optimizer.step()
                
                # Print losses
                if i % 100 == 0:
                    print("[Epoch %d/%d] [Batch %d/%d] Critic Loss: %.4f Generator Loss: %.4f" % 
                          (epoch, num_epochs, i, len(data_loader), critic_loss.item(), generator_loss.item()))

# Usage example
input_dim = 100
output_dim = 784

# Load your dataset and create a DataLoader
dataset = ...
data_loader = DataLoader(dataset, batch_size=64, shuffle=True)

wgan = WGAN(input_dim, output_dim)
wgan.train(data_loader, num_epochs=10)

请注意，这只是一个简单的WGAN实现示例，你可能需要根据你的具体任务进行适当的修改和调整。同时，确保你的数据集已经正确加载和预处理，并用适当的损失函数和优化器进行训练。