如何在每个批次中创建具有不同元素的火炬对角线矩阵？

import random

def create_torch_diagonal_matrix(batch_size, n, m):
    matrices = []
    for _ in range(batch_size):
        matrix = [[0] * n for _ in range(n)]
        random_array = [random.randint(1, m) for _ in range(n)]
        for i in range(n):
            for j in range(n):
                if i == j:
                    matrix[i][j] = random_array[i]
        matrices.append(matrix)
    return matrices

# 示例调用
batch_size = 3
n = 5
m = 10
matrices = create_torch_diagonal_matrix(batch_size, n, m)
for matrix in matrices:
    for row in matrix:
        print(row)
    print()