OBJECT_MEMORY_LIMIT_MB = 100


def detect_torch_usage(test_context: str) -> bool:
    """Detect usage of torch modules in code."""
    return "torch" in test_context or "Tensor" in test_context


def get_tensor_size_note(test_context: str) -> list[str]:
    """If torch detected in code, return a note to reduce tensor size."""
    if detect_torch_usage(test_context):
        return [f"Do not generate torch tensors with size greater than {OBJECT_MEMORY_LIMIT_MB}MB."]
    return []
