class_docstr: MNIST classification task.
class_docstr_tok: ["MNIST","classification","task","."]
code: def __init__(self, batch_size: int = 1024, test: bool = False):
      self.max_steps = 1
      self.obs_shape = tuple([28, 28, 1])
      self.act_shape = tuple(
          [
              10,
          ]
      )
  
      # Delayed importing of torchvision
  
      try:
          from torchvision import datasets
      except ModuleNotFoundError:
          print("You need to install torchvision for this task.")
          print("  pip install torchvision")
          sys.exit(1)
  
      dataset = datasets.MNIST("./data", train=not test, download=True)
      data = np.expand_dims(dataset.data.numpy() / 255.0, axis=-1)
      labels = dataset.targets.numpy()
  
      def reset_fn(key):
          if test:
              batch_data, batch_labels = data, labels
          else:
              batch_data, batch_labels = sample_batch(
                  key, data, labels, batch_size
              )
          return State(obs=batch_data, labels=batch_labels)
  
      self._reset_fn = jax.jit(jax.vmap(reset_fn))
  
      def step_fn(state, action):
          if test:
              reward = accuracy(action, state.labels)
          else:
              reward = -loss(action, state.labels)
          return state, reward, jnp.ones(())
  
      self._step_fn = jax.jit(jax.vmap(step_fn))
  
file_path: evojax/evojax/task/mnist.py
imports: import sys
  import numpy
  from typing import Tuple
  import jax
  import jax.numpy
  from jax import random
  from flax.struct import dataclass
  from evojax.task.base import VectorizedTask
  from evojax.task.base import TaskState
label_desc: Use image node_id 1 to create a new MNIST object from inherited base classes: VectorizedTask with example: obj = MNIST(batch_size, test)