import torch
from transformers import AutoConfig, AutoModel
from configuration_davit import DaViTConfig
from modeling_davit import DaViTModel

# Register the configuration and model
AutoConfig.register("davit", DaViTConfig)
AutoModel.register(DaViTConfig, DaViTModel)

# Step 1: Create a configuration object
config = DaViTConfig()

# Step 2: Create a model object
model = AutoModel.from_config(config)

# Step 3: Run a forward pass
# Generate a random sample input tensor with shape (batch_size, channels, height, width)
batch_size = 2
channels = 3
height = 224
width = 224
sample_input = torch.randn(batch_size, channels, height, width)

# Pass the sample input through the model
output = model(sample_input)

# Print the output shape
print(f"Output shape: {output.shape}")

# Expected output shape: (batch_size, projection_dim)