请帮我把这段代码改成def的形式data = torch.randn(16992, 307, 12, 2) week_feature = torch.zeros((data.shape[0], 1)) # 遍历所有时间点，获取对应的星期并将星期转化为对应的数字 for i in range(data.shape[0]): date_str = '2016-01-01 00:00:00' # 假设数据集中的时间从2016年1月1日开始 date = datetime.datetime.strptime(date_str, '%Y-%m-%d %H:%M:%S') date += datetime.timedelta(minutes=5*i) # 每个时间点间隔5分钟 week = date.weekday() # 获取星期，0表示星期一，1表示星期二，以此类推 week_feature[i] = week + 1 # 将星期转化为对应的数字并将其存入星期特征张量中 # 将星期特征添加到原有的数据集中 week_feature=week_feature.unsqueeze(1).repeat(1, 307, 12, 1) week_feature=week_feature.reshape(16992,307,12,1) data = torch.cat((data, week_feature), dim=-1) # week_feature = week_feature.repeat(1, 307, 12, 1) # 将星期特征张量与数据集在最后一个维度上拼接 # data = torch.cat((data, week_feature), dim=-1) # 输出结果 print(data) print(data.shape) ，pytorch版本

好的，可以将代码改成如下def函数形式：

import torch
import datetime

def add_week_feature(data):
    """
    将星期特征添加到原有的数据集中

    :param data: 原始数据集，大小为 (batch_size, seq_len, num_features)
    :return: 添加星期特征后的数据集，大小为 (batch_size, seq_len, num_features+1)
    """
    batch_size, seq_len, num_features = data.shape
    week_feature = torch.zeros((batch_size, seq_len, 1))

    # 遍历所有时间点，获取对应的星期并将星期转化为对应的数字
    for i in range(batch_size):
        date_str = '2016-01-01 00:00:00'  # 假设数据集中的时间从2016年1月1日开始
        date = datetime.datetime.strptime(date_str, '%Y-%m-%d %H:%M:%S')
        for j in range(seq_len):
            date += datetime.timedelta(minutes=5*j)  # 每个时间点间隔5分钟
            week = date.weekday()  # 获取星期，0表示星期一，1表示星期二，以此类推
            week_feature[i, j, 0] = week + 1  # 将星期转化为对应的数字并将其存入星期特征张量中

    # 将星期特征添加到原有的数据集中
    week_feature = week_feature.repeat(1, 1, num_features).reshape(batch_size, seq_len, num_features, 1)
    data = torch.cat((data, week_feature), dim=-1)

    return data

其中，`add_week_feature` 函数的输入参数 `data` 是原始数据集，大小为 `(batch_size, seq_len, num_features)`，输出为添加星期特征后的数据集，大小为 `(batch_size, seq_len, num_features+1)`。函数中的代码与原始代码功能相同，只是将循环中的索引变量 `i` 和 `j` 分别表示批次和时间序列。