code: def test_model_gpu():
      pytest.importorskip("ml_datasets")
      import ml_datasets
  
      with use_ops("cupy"):
          n_hidden = 32
          dropout = 0.2
          (train_X, train_Y), (dev_X, dev_Y) = ml_datasets.mnist()
          model = chain(
              Relu(nO=n_hidden, dropout=dropout),
              Relu(nO=n_hidden, dropout=dropout),
              Softmax(),
          )
          # make sure the data is on the right device
          train_X = model.ops.asarray(train_X)
          train_Y = model.ops.asarray(train_Y)
          dev_X = model.ops.asarray(dev_X)
          dev_Y = model.ops.asarray(dev_Y)
  
          model.initialize(X=train_X[:5], Y=train_Y[:5])
          optimizer = Adam(0.001)
          batch_size = 128
  
          for i in range(2):
              batches = model.ops.multibatch(
                  batch_size, train_X, train_Y, shuffle=True
              )
              for X, Y in batches:
                  Yh, backprop = model.begin_update(X)
                  backprop(Yh - Y)
                  model.finish_update(optimizer)
              # Evaluate and print progress
              correct = 0
              total = 0
              for X, Y in model.ops.multibatch(
                  batch_size, dev_X, dev_Y
              ):
                  Yh = model.predict(X)
                  correct += (
                      Yh.argmax(axis=1) == Y.argmax(axis=1)
                  ).sum()
                  total += Yh.shape[0]
  
imports: import threading
  import time
  from collections import Counter
  import numpy
  import pytest
  from thinc.api import Adam, CupyOps, Dropout, Linear, Model, Relu, Shim, Softmax, chain, change_attr_values, concatenate, set_dropout_rate, use_ops, with_debug, wrap_model_recursive
  from thinc.compat import has_cupy_gpu
  from ..util import make_tempdir
label_desc: Use image node_id 25 for calling a global function with example usage: test_model_gpu() without return types