Python 如何对 torch 张量中的数据进行切片

我已经把我的数据放到了 torch 张量上，现在我想分成大小为64的批次.我得到了以下代码:

batch = 0
BATCH_SIZE = 64
X_train = x_scaled.to(device)
y_train = y_scaled.to(device)

for batch in range(0,len(X_train[0]),BATCH_SIZE):
   ### Training
   model.train() # train mode is on by default after construction
   # 1. Forward pass
   y_pred = model(X_train[0:2][batch:batch+BATCH_SIZE])

张量的形状是:torch.Size([2, 11938]).我想把它切成[2，64].但是，它没有正确切片并给出错误:mat1 and mat2 shapes cannot be multiplied (2x11938 and 2x64)

我想要的是:

tensor([[0.0000, 0.0002, 0.0004, 0.0005, 0.0007, 0.0009, 0.0011, 0.0013, 0.0014,
    0.0016, 0.0018, 0.0018, 0.0020, 0.0022, 0.0023, 0.0025, 0.0027, 0.0029,
    0.0029, 0.0031, 0.0032, 0.0034, 0.0036, 0.0038, 0.0040, 0.0041, 0.0043,
    0.0045, 0.0047, 0.0049, 0.0051, 0.0052, 0.0054, 0.0056, 0.0058, 0.0060,
    0.0061, 0.0061, 0.0063, 0.0065, 0.0067, 0.0069, 0.0070, 0.0072, 0.0074,
    0.0076, 0.0078, 0.0079, 0.0081, 0.0083, 0.0083, 0.0085, 0.0087, 0.0088,
    0.0090, 0.0092, 0.0094, 0.0094, 0.0096, 0.0097, 0.0099, 0.0101, 0.0103,
    0.0105],[0.0684, 0.0684, 0.0684, 0.0684, 0.0684, 0.0684, 0.0684, 0.0684, 0.0684,
    0.0703, 0.0703, 0.0703, 0.0684, 0.0684, 0.0703, 0.0703, 0.0703, 0.0703,
    0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703,
    0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703,
    0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703,
    0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703, 0.0703,
    0.0703, 0.0703, 0.0712, 0.0712, 0.0712, 0.0712, 0.0712, 0.0712, 0.0712,
    0.0712]], device='cuda:0', dtype=torch.float64)

我得到的:

tensor([[0.0000e+00, 1.8038e-04, 3.6076e-04,  ..., 9.9964e-01, 9.9982e-01,
         1.0000e+00],
        [6.8395e-02, 6.8395e-02, 6.8395e-02,  ..., 5.7695e-01, 5.7695e-01,
         5.7695e-01]], device='cuda:0', dtype=torch.float64)

我怎样才能把 torch 张量切成所需的形状？