How to train a Keras model with a PyTorch DataLoader

Keras 3 training loops can receive batches from PyTorch data utilities instead of only arrays or TensorFlow datasets. That lets a project keep existing torch.utils.data.Dataset and DataLoader code while training the model with Keras compile() and fit().

A DataLoader is already batch-aware. It owns the batch size, sample order, shuffling, worker settings, and any dataset-side transforms, so model.fit() should receive the loader as x without a separate y array or batch_size argument.

Use the torch backend when the same process should keep tensors, loader batches, and model execution in the PyTorch runtime. Keras can consume PyTorch loaders with other backends too, but starting with KERAS_BACKEND=torch avoids extra framework conversion when the project data layer already uses PyTorch tensors.

1. Create train_with_torch_dataloader.py with backend selection and imports.

   train_with_torch_dataloader.py

   import os
    
   os.environ["KERAS_BACKEND"] = "torch"
    
   import keras
   import torch
   from torch.utils.data import DataLoader, TensorDataset
    
   keras.utils.set_random_seed(19)
   torch.manual_seed(19)

   Set KERAS_BACKEND before the first keras import. The Python environment must have both Keras and PyTorch installed.
   Related: How to set the Keras backend
   Related: How to install Keras with pip

2. Add feature and target tensors for a small binary training set.

   x_train = torch.tensor(
       [
           [0.00, 0.10, 0.15],
           [0.20, 0.35, 0.30],
           [0.40, 0.45, 0.55],
           [0.65, 0.70, 0.60],
           [0.75, 0.85, 0.90],
           [0.90, 0.95, 0.80],
           [0.10, 0.20, 0.25],
           [0.55, 0.60, 0.65],
       ],
       dtype=torch.float32,
   )
   y_train = torch.tensor(
       [[0.0], [0.0], [0.0], [1.0], [1.0], [1.0], [0.0], [1.0]],
       dtype=torch.float32,
   )

   The target tensor uses shape (samples, 1) because the model returns one sigmoid score for each sample.

3. Wrap the tensors in a TensorDataset and DataLoader.

   train_dataset = TensorDataset(x_train, y_train)
   train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)

   shuffle=True belongs to the DataLoader. Keras receives complete batches from the loader instead of choosing the sample order itself.

4. Inspect one loader batch before building the model.

   batch_x, batch_y = next(iter(train_loader))
    
   print(f"backend: {keras.config.backend()}")
   print(f"loader batches: {len(train_loader)}")
   print(f"batch x shape: {tuple(batch_x.shape)}")
   print(f"batch y shape: {tuple(batch_y.shape)}")
   print(f"batch x dtype: {batch_x.dtype}")
   print(f"batch y dtype: {batch_y.dtype}")

5. Add a Sequential model that accepts three input features.

   model = keras.Sequential(
       [
           keras.layers.Input(shape=(3,)),
           keras.layers.Dense(8, activation="relu"),
           keras.layers.Dense(1, activation="sigmoid"),
       ]
   )

   Related: How to create a Sequential model in Keras

6. Compile the model with a binary loss and accuracy metric.

   model.compile(
       optimizer=keras.optimizers.Adam(learning_rate=0.05),
       loss=keras.losses.BinaryCrossentropy(),
       metrics=[keras.metrics.BinaryAccuracy(name="accuracy")],
   )

   Related: How to compile a model in Keras

7. Train the model from the DataLoader.

   history = model.fit(train_loader, epochs=3, shuffle=False, verbose=2)

   Do not pass a separate batch_size or y value when x is a DataLoader. Passing shuffle=False keeps Keras from applying its own shuffle flag to a loader that already controls sample order.

8. Evaluate the trained model and predict on the inspected batch.

   metrics = model.evaluate(train_loader, verbose=0, return_dict=True)
   predictions = model.predict(batch_x, verbose=0)

9. Print the training and smoke-test evidence.

   print("history keys:", ", ".join(sorted(history.history.keys())))
   print(f"epochs completed: {len(history.history['loss'])}")
   print(f"final accuracy: {history.history['accuracy'][-1]:.4f}")
   print("evaluation keys:", ", ".join(sorted(metrics.keys())))
   print(f"prediction shape: {predictions.shape}")

10. Run the script and confirm that fit() consumes the DataLoader batches.

    $ python train_with_torch_dataloader.py
    backend: torch
    loader batches: 2
    batch x shape: (4, 3)
    batch y shape: (4, 1)
    batch x dtype: torch.float32
    batch y dtype: torch.float32
    Epoch 1/3
    2/2 - 0s - 17ms/step - accuracy: 0.2500 - loss: 0.7044
    Epoch 2/3
    2/2 - 0s - 5ms/step - accuracy: 0.5000 - loss: 0.6541
    Epoch 3/3
    2/2 - 0s - 6ms/step - accuracy: 0.5000 - loss: 0.6378
    history keys: accuracy, loss
    epochs completed: 3
    final accuracy: 0.5000
    evaluation keys: accuracy, loss
    prediction shape: (4, 1)