utils/initialize_input.py

import torch
from typing import Tuple
def initialize_input(shape:Tuple[int, ...]=(1,3,224,224),
                     input_type:str = "image",
                     seed:int=42) -> torch.Tensor:
  """
    Initialize input data as random noise.
    Args:
        shape (tuple): Shape of the input tensor (aligned with the model) .
        type (str): catagorizes as image/text for reference (feel free to leave as "")
        seed (int): random seed
    Returns:
        torch.Tensor: Randomly initialized input tensor.
  """

  torch.manual_seed(seed)
  torch.cuda.manual_seed(seed)

  return torch.randn(shape, requires_grad=True), input_type