text: # Copyright 2021 The HuggingFace Team. All rights reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import numpy as np
  import torch
  
  
  class RegressionDataset:
      def __init__(self, a=2, b=3, length=64, seed=None):
          if seed is not None:
              np.random.seed(seed)
          self.length = length
          self.x = np.random.normal(size=(length,)).astype(np.float32)
          self.y = a * self.x + b + np.random.normal(scale=0.1, size=(length,)).astype(np.float32)
  
      def __len__(self):
          return self.length
  
      def __getitem__(self, i):
          return {"x": self.x[i], "y": self.y[i]}
  
  
  class RegressionModel(torch.nn.Module):
      def __init__(self, a=0, b=0, double_output=False):
          super().__init__()
          self.a = torch.nn.Parameter(torch.tensor(a).float())
          self.b = torch.nn.Parameter(torch.tensor(b).float())
          self.first_batch = True
  
      def forward(self, x=None):
          if self.first_batch:
              print(f"Model dtype: {self.a.dtype}, {self.b.dtype}. Input dtype: {x.dtype}")
              self.first_batch = False
          return x * self.a + self.b
  