docstr: :param batch_size: number of samples to return
  :param use_gpu: whether to put samples on gpu
  State's shape is SEQ_LEN x BATCH_SIZE x STATE_DIM, for example.
  By default, MDN-RNN consumes data with SEQ_LEN as the first dimension.
docstr_tok: [":","param","batch_size",":","number","of","samples","to","return",":","param","use_gpu",":","whether","to","put","samples","on","gpu","State","'s","shape","is","SEQ_LEN","x","BATCH_SIZE","x","STATE_DIM",",","for","example",".","By","default",",","MDN-RNN","consumes","data","with","SEQ_LEN","as","the","first","dimension","."]
code: def sample_memories(
      self, batch_size, use_gpu=False
  ) -> rlt.MemoryNetworkInput:
      """
      :param batch_size: number of samples to return
      :param use_gpu: whether to put samples on gpu
      State's shape is SEQ_LEN x BATCH_SIZE x STATE_DIM, for example.
      By default, MDN-RNN consumes data with SEQ_LEN as the first dimension.
      """
      sample_indices = np.random.randint(
          self.memory_size, size=batch_size
      )
      device = torch.device("cuda") if use_gpu else torch.device("cpu")
      # state/next state shape: batch_size x seq_len x state_dim
      # action shape: batch_size x seq_len x action_dim
      # reward/not_terminal shape: batch_size x seq_len
      state, action, next_state, reward, not_terminal = map(
          lambda x: stack(x).float().to(device),
          zip(*self.deque_sample(sample_indices)),
      )
  
      # make shapes seq_len x batch_size x feature_dim
      state, action, next_state, reward, not_terminal = transpose(
          state, action, next_state, reward, not_terminal
      )
  
      return rlt.MemoryNetworkInput(
          state=rlt.FeatureData(float_features=state),
          reward=reward,
          time_diff=torch.ones_like(reward).float(),
          action=rlt.FeatureData(float_features=action),
          next_state=rlt.FeatureData(float_features=next_state),
          not_terminal=not_terminal,
          step=None,
      )
  
imports: import logging
  from collections import deque
  from typing import NamedTuple
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from reagent.core import types
  from reagent.core.torch_utils import stack
  from torch.distributions.normal import Normal
label_desc: Use image node_id 3 for calling the MDNRNNMemoryPool obj's underlying member method code with example usage: obj.sample_memories(batch_size, use_gpu) and returns: rlt