text: import pytest
  import torch
  
  from wild_relation_network import RelationNetwork
  
  
  @pytest.mark.parametrize('use_object_triples', [False, True])
  @pytest.mark.parametrize('use_layer_norm', [False, True])
  def test_forward(use_object_triples, use_layer_norm):
      x = torch.rand(4, 8, 64)
      rn = RelationNetwork(
          num_objects=8,
          object_size=64,
          out_size=32,
          use_object_triples=use_object_triples,
          use_layer_norm=use_layer_norm
      )
      logits = rn(x)
      assert logits.shape == (4, 32)
      assert logits.isfinite().all()
  