MxNet Gluon, правильный способ обучения модели Sequential() на графическом процессоре

Question

Я пытался обучить модель с использованием графического процессора на сервере, но получаю сообщение об ошибке:

      Check failed: e == CUBLAS_STATUS_SUCCESS (13 vs. 0) : cuBLAS: CUBLAS_STATUS_EXECUTION_FAILED

Я нашел несколько похожих тем на форуме, и все они, похоже, указывают на некоторые проблемы с установленными версиями CUDA. Однако на той же машине я могу обучить некоторые модели обнаружения объектов на графическом процессоре в той же среде без каких-либо ошибок, поэтому я предполагаю, что проблема в моем коде.

Это то, что я делаю до сих пор:

      import mxnet as mx
from mxnet import autograd, gluon, init
from mxnet import ndarray as nd
import numpy as np
from mxnet.gluon import nn
from mxnet import init
from mxnet.gluon import loss as gloss
import pandas as pd


def main()

    ## read, preprocess and split data
    df_data = pd.read_csv('some_file.csv')
    df_data = pre_process(df_data)
    X_train, y_train, X_test, y_test = split_data(df_data)


    train(X_train, X_test, y_train, y_test, lr, batch_size, nr_epochs)


def train(X_train, X_test, y_train, y_test, lr, batch_size, nr_epochs):
    ctx = mx.gpu(3)
    y_train = mx.nd.array(y_train.to_numpy().reshape(-1,1), dtype=np.float32, ctx=ctx)
    y_test = mx.nd.array(y_test.to_numpy().reshape(-1,1), dtype=np.float32, ctx=ctx)
    X_train = mx.nd.array(X_train.to_numpy(), dtype=np.float32, ctx=ctx)
    X_test = mx.nd.array(X_test.to_numpy(), dtype=np.float32, ctx=ctx)

    ##--------------------
    ##   building model
    ##--------------------
    batch = batch_size
    epochs = nr_epochs
    dataset = gluon.data.dataset.ArrayDataset(X_train, y_train)
    data_loader = gluon.data.DataLoader(dataset, batch_size=batch, shuffle=True)

    model = nn.Sequential()
    model.add(nn.Dense(64, activation='relu'))
    model.add(nn.Dense(1))
    model.initialize(init.Normal(sigma=0.01), ctx)
    model.collect_params().reset_ctx(ctx)
    loss = gloss.L2Loss()
    trainer = gluon.Trainer(model.collect_params(), 'sgd', {'learning_rate': lr})

    ##--------------------
    ##   training
    ##--------------------
    for epoch in range(1, epochs + 1):
        for X_batch, Y_batch in data_loader:
            with autograd.record():
                l = loss(model(X_batch), Y_batch)
            l.backward()
            trainer.step(batch)

ошибка возникает, даже если я пытаюсь сделать что-то простое, например:

      print(l)

внутри цикла for, поэтому я предполагаю, что здесь что-то не так?

Я использую:
mxnet-cu90mkl=1.5.1.post0