Fix translate fill

ffcv/pipeline/graph.py

@@ -2,7 +2,12 @@
 import warnings
 import ast
 
-import astor
+try:
+    # Useful for debugging
+    import astor
+except ImportError:
+    pass
+
 from collections import defaultdict
 from typing import Callable, Dict, List, Optional, Sequence, Set
 from abc import ABC, abstractmethod

ffcv/transforms/ops.py

@@ -42,6 +42,8 @@ class ToDevice(Operation):
     def __init__(self, device, non_blocking=True):
         super().__init__()
         self.device = device
+        # assert isinstance(device, ch.device), \
+        #    f'Make sure device is a ch.device (not a {type(device)})'
         self.non_blocking = non_blocking
 
     def generate_code(self) -> Callable:
@@ -155,4 +157,4 @@ def convert(inp, dst):
         return convert
 
     def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]:
-        return replace(previous_state, dtype=self.target_dtype, jit_mode=False), None
+        return replace(previous_state, dtype=self.target_dtype, jit_mode=False), None

ffcv/transforms/translate.py

@@ -30,18 +30,15 @@ def __init__(self, padding: int, fill: Tuple[int, int, int] = (0, 0, 0)):
     def generate_code(self) -> Callable:
         my_range = Compiler.get_iterator()
         pad = self.padding
+        fill = self.fill
 
         def translate(images, dst):
             n, h, w, _ = images.shape
-            # y_coords = randint(low=0, high=2 * pad + 1, size=(n,))
-            # x_coords = randint(low=0, high=2 * pad + 1, size=(n,))
-            # dst = fill
-
+            dst[:] = fill
             dst[:, pad:pad+h, pad:pad+w] = images
             for i in my_range(n):
                 y_coord = randint(low=0, high=2 * pad + 1)
                 x_coord = randint(low=0, high=2 * pad + 1)
-                # images[i] = dst[i, y_coords[i]:y_coords[i]+h, x_coords[i]:x_coords[i]+w]
                 images[i] = dst[i, y_coord:y_coord+h, x_coord:x_coord+w]
 
             return images