code: def _test_forward(self, gpu):
      in_size = 2
      out0_size = 3
      out1_size = 4
      out2_size = 1
  
      par = RecurrentBranched(
          nn.LSTM(
              num_layers=1, input_size=in_size, hidden_size=out0_size
          ),
          RecurrentSequential(
              nn.RNN(
                  num_layers=1,
                  input_size=in_size,
                  hidden_size=out1_size,
              ),
          ),
          RecurrentSequential(
              nn.Linear(in_size, out2_size),
          ),
      )
  
      if gpu >= 0:
          device = torch.device("cuda:{}".format(gpu))
          par.to(device)
      else:
          device = torch.device("cpu")
  
      seqs_x = [
          torch.rand(1, in_size, device=device),
          torch.rand(3, in_size, device=device),
      ]
  
      packed_x = nn.utils.rnn.pack_sequence(
          seqs_x, enforce_sorted=False
      )
  
      # Concatenated output should be a tuple of three variables.
      out, rs = par(packed_x, None)
      self.assertIsInstance(out, tuple)
      self.assertEqual(len(out), len(par))
      self.assertEqual(out[0].data.shape, (4, out0_size))
      self.assertEqual(out[1].data.shape, (4, out1_size))
      self.assertEqual(out[2].data.shape, (4, out2_size))
  
      self.assertIsInstance(rs, tuple)
      self.assertEqual(len(rs), len(par))
  
      # LSTM
      self.assertIsInstance(rs[0], tuple)
      self.assertEqual(len(rs[0]), 2)
      self.assertEqual(rs[0][0].shape, (1, len(seqs_x), out0_size))
      self.assertEqual(rs[0][1].shape, (1, len(seqs_x), out0_size))
  
      # RecurrentSequential(RNN)
      self.assertIsInstance(rs[1], tuple)
      self.assertEqual(len(rs[1]), 1)
      self.assertEqual(rs[1][0].shape, (1, len(seqs_x), out1_size))
  
      # RecurrentSequential(Linear)
      self.assertIsInstance(rs[2], tuple)
      self.assertEqual(len(rs[2]), 0)
  
file_path: pfrl/tests/nn_tests/test_recurrent_branched.py
imports: import unittest
  import pytest
  import torch
  from torch import nn
  from pfrl.nn import RecurrentBranched, RecurrentSequential
  from pfrl.testing import torch_assert_allclose
  from pfrl.utils.recurrent import concatenate_recurrent_states, get_recurrent_state_at, mask_recurrent_state_at, one_step_forward
label_desc: Use image node_id 1 for calling the TestRecurrentBranched obj's underlying member method code with example usage: obj._test_forward(gpu) without return types